Memory problem fixed by Claude

[benikm91]

From Claude:

Fix Python reference leak in pytree element access

ScalaPy 0.5.3's bracketAccess leaks one Python reference per call: CPythonInterpreter.selectBracket wraps PyObject_GetItem — which returns a new reference — in PyValue.fromBorrowed, which takes a second one. Only one 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. (Attribute access via select correctly uses fromNew, and PyValue.getTuple correctly pairs the borrowed PyTuple_GetItem with fromBorrowed — selectBracket is the odd one out.)

TensorTree.fromPyTree reads every leaf of a returned pytree with bracketAccess, so each jitted call leaked its entire result tree. In a DETR training loop that was ~556 Python objects per step, growing without bound. The visible symptom was not memory but speed: a training loop that calls dimwit.gc() per step pays O(tracked objects) for gc.collect(), so step time grew linearly — throughput halved over 1900 steps and kept falling.

Add dimwit.python.PyIndex.itemAt, which indexes via __getitem__ and so goes through the ordinary attribute-call path (PyObject_GetAttrString + PyValue.fromNew), where refcounting is balanced. Replace bracketAccess with it in TensorTree, Jit, Autodiff, LinearAlgebra and StructuralOps.

Measured on a DETR training loop, 1500 steps:

before: 450 -> 207 samples/sec and still falling,
live JAX arrays +1/step, Python objects +556/step
after: ~470 samples/sec stable,
live JAX arrays constant at 368, Python objects constant

Removing the now-unnecessary per-step dimwit.gc() on top of this takes the same loop to ~950-1000 samples/sec.

[Copilot]

Pull request overview

This PR addresses a Python reference leak triggered by ScalaPy’s bracketAccess when indexing Python objects returned from JAX (notably pytrees), which could cause unbounded growth in live Python/JAX objects and degrade performance over long training loops.

Changes:

• Introduces dimwit.python.PyIndex.itemAt as a leak-free indexing helper that routes access via __getitem__.
• Replaces bracketAccess with itemAt across TensorTree pytree extraction and key JAX integration points (Jit donating reducers, Autodiff value-and-grad, linear algebra returns, and structural indexing).

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
core/src/main/scala/dimwit/tensortree/TensorTree.scala	Switches pytree leaf and tuple/list extraction from bracketAccess to leak-free itemAt.
core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala	Uses itemAt for JAX tensor indexing and .at[...] updates to avoid leaking references.
core/src/main/scala/dimwit/python/PyIndex.scala	Adds the itemAt extension method(s) on py.Dynamic as a replacement for bracketAccess.
core/src/main/scala/dimwit/linalg/LinearAlgebra.scala	Replaces tuple-style result indexing (qr/eigh/svd) with itemAt.
core/src/main/scala/dimwit/jax/Jit.scala	Uses itemAt when unpacking multi-result returns from jitted donating reducers.
core/src/main/scala/dimwit/autodiff/Autodiff.scala	Uses itemAt when unpacking (value, grad) from the Python helper call.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[marcelluethi]

Thanks for the fix and also for the careful experiments.
This is great news that we have this leak fixed - and even get a speed improvement.

We should maybe open an issue in ScalaPy. While ScalaPy does not seem to be actively maintained, it would at least be a warning for others.

[marcelluethi]

I find this variant even more readable then bracketAccess

[benikm91]

FYI: I will test if I can remove the extra memory steps in GPT before merging this