diff --git a/core/src/main/scala/dimwit/autodiff/Autodiff.scala b/core/src/main/scala/dimwit/autodiff/Autodiff.scala index 613e471..43bd7a4 100644 --- a/core/src/main/scala/dimwit/autodiff/Autodiff.scala +++ b/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 diff --git a/core/src/main/scala/dimwit/jax/Jit.scala b/core/src/main/scala/dimwit/jax/Jit.scala index 5f39dbc..aed4573 100644 --- a/core/src/main/scala/dimwit/jax/Jit.scala +++ b/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) --- diff --git a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala b/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala index 824836a..f7334cb 100644 --- a/core/src/main/scala/dimwit/linalg/LinearAlgebra.scala +++ b/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. diff --git a/core/src/main/scala/dimwit/python/PyIndex.scala b/core/src/main/scala/dimwit/python/PyIndex.scala new file mode 100644 index 0000000..eb12e07 --- /dev/null +++ b/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) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index a88d305..6bd88f3 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/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 diff --git a/core/src/main/scala/dimwit/tensortree/TensorTree.scala b/core/src/main/scala/dimwit/tensortree/TensorTree.scala index fb11cb4..c38c3fd 100644 --- a/core/src/main/scala/dimwit/tensortree/TensorTree.scala +++ b/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))