Memory problem fixed by Claude

core/src/main/scala/dimwit/autodiff/Autodiff.scala

@@ -1,5 +1,6 @@
 package dimwit.autodiff
 
+import dimwit.python.PyIndex.itemAt
 import dimwit.OnError
 import dimwit.jax.Jax
 import dimwit.tensor.Tensor
@@ -68,8 +69,8 @@ object Autodiff:
     (params: Input) =>
       val pyParams = inTree.toPyTree(params)
       val r = gpy(pyParams)
-      val pyValue = r.bracketAccess(0)
-      val pyGrad = r.bracketAccess(1)
+      val pyValue = r.itemAt(0)
+      val pyGrad = r.itemAt(1)
       (Tensor(pyValue), Grad(inTree.fromPyTree(pyGrad).asInstanceOf[Input]))
 
   def jacobian[In, Out](f: In => Out)(using

core/src/main/scala/dimwit/jax/Jit.scala

@@ -1,5 +1,6 @@
 package dimwit.jax
 
+import dimwit.python.PyIndex.itemAt
 import dimwit.OnError
 import dimwit.tensortree.TensorTree
 import dimwit.jax.Jax
@@ -271,7 +272,7 @@ object JitDonating:
     val jitted = pyJit(fpy, Map("donate_argnums" -> Tuple2(0, 1)))
     def apply(r1: Donatable, r2: Donatable2): (Donatable, Donatable2) =
       val res = jitted(r1, r2).as[Jax.PyDynamic]
-      (res.bracketAccess(0), res.bracketAccess(1))
+      (res.itemAt(0), res.itemAt(1))
 
   // Three Params
 
@@ -301,7 +302,7 @@ object JitDonating:
     def apply(t1: T1, r1: Donatable, r2: Donatable2): (Donatable, Donatable2) =
       val pyT1 = TensorTree[T1].toPyTree(t1)
       val res = jitted(pyT1, r1, r2).as[Jax.PyDynamic]
-      (res.bracketAccess(0), res.bracketAccess(1))
+      (res.itemAt(0), res.itemAt(1))
 
   case class JitReducerI0O3[R1: TensorTree, R2: TensorTree, R3: TensorTree](f: (R1, R2, R3) => (R1, R2, R3)) extends JitReducerO3[R1, R2, R3]:
     val fpy = (r1: Donatable, r2: Donatable2, r3: Donatable3) =>
@@ -314,7 +315,7 @@ object JitDonating:
     val jitted = pyJit(fpy, Map("donate_argnums" -> Tuple3(0, 1, 2)))
     def apply(r1: Donatable, r2: Donatable2, r3: Donatable3): (Donatable, Donatable2, Donatable3) =
       val res = jitted(r1, r2, r3).as[Jax.PyDynamic]
-      (res.bracketAccess(0), res.bracketAccess(1), res.bracketAccess(2))
+      (res.itemAt(0), res.itemAt(1), res.itemAt(2))
 
   // Four Params
 
@@ -348,7 +349,7 @@ object JitDonating:
       val pyT1 = TensorTree[T1].toPyTree(t1)
       val pyT2 = TensorTree[T2].toPyTree(t2)
       val res = jitted(pyT1, pyT2, r1, r2).as[Jax.PyDynamic]
-      (res.bracketAccess(0), res.bracketAccess(1))
+      (res.itemAt(0), res.itemAt(1))
 
   case class JitReducerI1O3[R1: TensorTree, R2: TensorTree, R3: TensorTree, T1: TensorTree](f: (T1, R1, R2, R3) => (R1, R2, R3)) extends JitReducerO3[R1, R2, R3]:
     val fpy = (t1: Jax.PyDynamic, r1: Donatable, r2: Donatable2, r3: Donatable3) =>
@@ -363,7 +364,7 @@ object JitDonating:
     def apply(t1: T1, r1: Donatable, r2: Donatable2, r3: Donatable3): (Donatable, Donatable2, Donatable3) =
       val pyT1 = TensorTree[T1].toPyTree(t1)
       val res = jitted(pyT1, r1, r2, r3).as[Jax.PyDynamic]
-      (res.bracketAccess(0), res.bracketAccess(1), res.bracketAccess(2))
+      (res.itemAt(0), res.itemAt(1), res.itemAt(2))
 
   case class JitReducerI0R4[R1: TensorTree, R2: TensorTree, R3: TensorTree, R4: TensorTree](f: (R1, R2, R3, R4) => (R1, R2, R3, R4)) extends JitReducerO4[R1, R2, R3, R4]:
     val fpy = (r1: Donatable, r2: Donatable2, r3: Donatable3, r4: Donatable4) =>
@@ -377,7 +378,7 @@ object JitDonating:
     val jitted = pyJit(fpy, Map("donate_argnums" -> Tuple4(0, 1, 2, 3)))
     def apply(r1: Donatable, r2: Donatable2, r3: Donatable3, r4: Donatable4): (Donatable, Donatable2, Donatable3, Donatable4) =
       val res = jitted(r1, r2, r3, r4).as[Jax.PyDynamic]
-      (res.bracketAccess(0), res.bracketAccess(1), res.bracketAccess(2), res.bracketAccess(3))
+      (res.itemAt(0), res.itemAt(1), res.itemAt(2), res.itemAt(3))
 
   // --- Helper Methods (Constructors) ---
 

core/src/main/scala/dimwit/linalg/LinearAlgebra.scala

@@ -1,5 +1,6 @@
 package dimwit.linalg
 
+import dimwit.python.PyIndex.itemAt
 import dimwit.jax.Jax
 import dimwit.tensor.Axis
 import dimwit.tensor.Label
@@ -129,7 +130,7 @@ object LinearAlgebra:
         case QRMode.Reduced  => "reduced"
         case QRMode.Complete => "complete"
     )
-    (q = Tensor[(LRow, LBasis), V](qr.bracketAccess(0)), r = Tensor[(LBasis, LCol), V](qr.bracketAccess(1)))
+    (q = Tensor[(LRow, LBasis), V](qr.itemAt(0)), r = Tensor[(LBasis, LCol), V](qr.itemAt(1)))
 
   /** Computes the eigenvalues and eigenvectors of a symmetric matrix `t`.
     * @param t The input tensor representing a symmetric matrix.
@@ -145,8 +146,8 @@ object LinearAlgebra:
       : (eigenvalues: Tensor1[LEig, V], eigenvectors: Tensor2[LSpace, LEig, V]) =
 
     val ret = Jax.jnp.linalg.eigh(t.jaxValue, UPLO = if upper then "U" else "L", symmetrize_input = symmetrizeInput)
-    val eigenvalues: Tensor1[LEig, V] = Tensor(ret.bracketAccess(0))
-    val eigenvectors: Tensor2[LSpace, LEig, V] = Tensor(ret.bracketAccess(1))
+    val eigenvalues: Tensor1[LEig, V] = Tensor(ret.itemAt(0))
+    val eigenvectors: Tensor2[LSpace, LEig, V] = Tensor(ret.itemAt(1))
     (eigenvalues = eigenvalues, eigenvectors = eigenvectors)
 
   /** Computes the singular value decomposition (SVD) of the tensor `t`.
@@ -164,9 +165,9 @@ object LinearAlgebra:
       : (U: Tensor2[LRow, LBasis, V], S: Tensor1[LSing, V], Vh: Tensor2[LBasis, LCol, V]) =
 
     val ret = Jax.jnp.linalg.svd(t.jaxValue, full_matrices = fullMatrices, hermitian = hermitian)
-    val u: Tensor2[LRow, LBasis, V] = Tensor(ret.bracketAccess(0))
-    val s: Tensor1[LSing, V] = Tensor(ret.bracketAccess(1))
-    val vh: Tensor2[LBasis, LCol, V] = Tensor(ret.bracketAccess(2))
+    val u: Tensor2[LRow, LBasis, V] = Tensor(ret.itemAt(0))
+    val s: Tensor1[LSing, V] = Tensor(ret.itemAt(1))
+    val vh: Tensor2[LBasis, LCol, V] = Tensor(ret.itemAt(2))
     (U = u, S = s, Vh = vh)
 
   /** Solves the linear equation Ax = b for x, where A is a square matrix and b is a vector.

core/src/main/scala/dimwit/python/PyIndex.scala

@@ -0,0 +1,31 @@
+package dimwit.python
+
+import me.shadaj.scalapy.py
+
+/** Leak-free replacement for ScalaPy's `bracketAccess`.
+  *
+  * `py.Dynamic.bracketAccess` goes through `CPythonInterpreter.selectBracket`,
+  * which wraps `PyObject_GetItem` — a function that returns a *new* reference —
+  * in `PyValue.fromBorrowed`, and that takes a second reference. Only one of the
+  * two is ever released, so every element read this way keeps its Python object
+  * alive for the rest of the process; for a JAX array that also pins its device
+  * buffer.
+  *
+  * Reading through `__getitem__` uses the ordinary attribute-call path
+  * (`PyObject_GetAttrString` + `PyValue.fromNew`), whose refcounting is
+  * balanced, so the element is released as soon as Scala drops it.
+  *
+  * This matters most in the training loop: every leaf of a jitted function's
+  * result is read out of the returned pytree, so a leak here grows with the
+  * number of steps and makes each `gc.collect()` progressively slower.
+  */
+private[dimwit] object PyIndex:
+
+  extension (value: py.Dynamic)
+    /** `value[index]`, without leaking a reference to the element. */
+    def itemAt(index: Int): py.Dynamic =
+      value.applyDynamic("__getitem__")(index)
+
+    /** `value[key]`, without leaking a reference to the element. */
+    def itemAt(key: py.Any): py.Dynamic =
+      value.applyDynamic("__getitem__")(key)

core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala

@@ -1,5 +1,6 @@
 package dimwit.tensor.tensorops
 
+import dimwit.python.PyIndex.itemAt
 import dimwit.jax.Einops
 import dimwit.jax.Jax
 import dimwit.tensor.Axis
@@ -581,7 +582,7 @@ object StructuralOps:
         labels: Labels[ev.RemainingAxes]
     ): Tensor[ev.RemainingAxes, V] =
       val pyIndices = tensor.calcPyIndices(inputs, ev.indices)
-      Tensor(tensor.jaxValue.bracketAccess(pyIndices))
+      Tensor(tensor.jaxValue.itemAt(pyIndices))
 
     /** Slice the given tensor, specifying the axis and index to slice at.
       *
@@ -667,7 +668,7 @@ object StructuralOps:
         labels: Labels[T]
     )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] =
       val pyIndices = tensor.calcPyIndices(inputs, ev.indices)
-      val result = tensor.jaxValue.at.bracketAccess(pyIndices).set(value.jaxValue)
+      val result = tensor.jaxValue.at.itemAt(pyIndices).set(value.jaxValue)
       Tensor(result)
 
     // Convenience overload for Float
@@ -679,7 +680,7 @@ object StructuralOps:
         labels: Labels[T]
     )(value: Float): Tensor[T, V] =
       val pyIndices = tensor.calcPyIndices(inputs, ev.indices)
-      val result = tensor.jaxValue.at.bracketAccess(pyIndices).set(value)
+      val result = tensor.jaxValue.at.itemAt(pyIndices).set(value)
       Tensor(result)
 
     // Convenience overload for AxisAtIndex

core/src/main/scala/dimwit/tensortree/TensorTree.scala

@@ -1,5 +1,6 @@
 package dimwit.tensortree
 
+import dimwit.python.PyIndex.itemAt
 import dimwit.jax.Jax
 import dimwit.tensor.*
 import dimwit.tensor.DType.Float32
@@ -95,9 +96,9 @@ object TensorTree: // extends TensorTreeLowPriority:
   ): (flatten: P => Tensor1[L, Float32], unflatten: Tensor1[L, Float32] => P) =
     val flattenUtil = py.module("jax.flatten_util")
     val result = flattenUtil.ravel_pytree(tt.toPyTree(reference)).as[py.Dynamic]
-    val unflattenPy = result.bracketAccess(1).as[py.Dynamic]
+    val unflattenPy = result.itemAt(1)
     val flatten = (p: P) =>
-      flatTree.fromPyTree(flattenUtil.ravel_pytree(tt.toPyTree(p)).as[py.Dynamic].bracketAccess(0))
+      flatTree.fromPyTree(flattenUtil.ravel_pytree(tt.toPyTree(p)).as[py.Dynamic].itemAt(0))
     val unflatten = (v: Tensor1[L, Float32]) => tt.fromPyTree(unflattenPy(flatTree.toPyTree(v)))
     (flatten = flatten, unflatten = unflatten)
 
@@ -180,14 +181,14 @@ object TensorTree: // extends TensorTreeLowPriority:
 
     def fromPyTree(pyVal: Jax.PyAny): (P1, P2) =
       val pyTuple = pyVal.as[py.Dynamic]
-      (t1.fromPyTree(pyTuple.bracketAccess(0)), t2.fromPyTree(pyTuple.bracketAccess(1)))
+      (t1.fromPyTree(pyTuple.itemAt(0)), t2.fromPyTree(pyTuple.itemAt(1)))
 
     def toNumpyTree(p: (P1, P2)): Jax.PyAny =
       py.Dynamic.global.tuple(Seq(t1.toNumpyTree(p._1), t2.toNumpyTree(p._2)).toPythonProxy)
 
     def fromNumpyTree(pyVal: Jax.PyAny): (P1, P2) =
       val pyTuple = pyVal.as[py.Dynamic]
-      (t1.fromNumpyTree(pyTuple.bracketAccess(0)), t2.fromNumpyTree(pyTuple.bracketAccess(1)))
+      (t1.fromNumpyTree(pyTuple.itemAt(0)), t2.fromNumpyTree(pyTuple.itemAt(1)))
 
   /** Instance for a list of tensor trees
     */
@@ -221,7 +222,7 @@ object TensorTree: // extends TensorTreeLowPriority:
     def fromPyTree(pyVal: Jax.PyAny): List[P] =
       val pyList = pyVal.as[py.Dynamic]
       val len = py.Dynamic.global.len(pyList).as[Int]
-      List.tabulate(len)(i => tp.fromPyTree(pyList.bracketAccess(i)))
+      List.tabulate(len)(i => tp.fromPyTree(pyList.itemAt(i)))
 
     def toNumpyTree(l: List[P]): Jax.PyAny =
       val pyItems = l.map(a => tp.toNumpyTree(a))
@@ -230,7 +231,7 @@ object TensorTree: // extends TensorTreeLowPriority:
     def fromNumpyTree(pyVal: Jax.PyAny): List[P] =
       val pyList = pyVal.as[py.Dynamic]
       val len = py.Dynamic.global.len(pyList).as[Int]
-      List.tabulate(len)(i => tp.fromNumpyTree(pyList.bracketAccess(i)))
+      List.tabulate(len)(i => tp.fromNumpyTree(pyList.itemAt(i)))
 
   given namedTupleInstance[N <: Tuple, V <: Tuple](using tt: TensorTree[V]): TensorTree[NamedTuple[N, V]] with
     def map(p: NamedTuple[N, V], f: [T <: Tuple, V2] => (Labels[T]) ?=> (Tensor[T, V2] => Tensor[T, V2])): NamedTuple[N, V] =
@@ -334,7 +335,7 @@ object TensorTree: // extends TensorTreeLowPriority:
     def fromPyTree(pyVal: Jax.PyAny): P =
       val pyTuple = pyVal.as[py.Dynamic]
       val elems = instances.zipWithIndex.map: (tc, index) =>
-        tc.fromPyTree(pyTuple.bracketAccess(index))
+        tc.fromPyTree(pyTuple.itemAt(index))
       m.fromProduct(Tuple.fromArray(elems.map(_.asInstanceOf[Object]).toArray))
 
     def toNumpyTree(p: P): Jax.PyAny =
@@ -345,5 +346,5 @@ object TensorTree: // extends TensorTreeLowPriority:
     def fromNumpyTree(pyVal: Jax.PyAny): P =
       val pyTuple = pyVal.as[py.Dynamic]
       val elems = instances.zipWithIndex.map: (tc, index) =>
-        tc.fromNumpyTree(pyTuple.bracketAccess(index))
+        tc.fromNumpyTree(pyTuple.itemAt(index))
       m.fromProduct(Tuple.fromArray(elems.map(_.asInstanceOf[Object]).toArray))