Repository: google/jax

Issue:
`jax.experimental.attrs`: Spurious "AssertionError: a jaxpr variable must be created only once per tracer"
### Description

The following code results in an erroneous jax internal assertion

> `AssertionError: a jaxpr variable must be created only once per tracer`

The issue is a use-after-free due to `experimental.attrs` creating a Tracer but not adequately ensuring its lifetime.

```python
import jax

from jax.experimental.attrs import jax_getattr, jax_setattr


class StatefulRNG:
    key: jax.Array

    def __init__(self, key: jax.Array):
        self.key = key

    def split(self) -> jax.Array:
        key = jax_getattr(self, "key")
        new_key, returned_key = jax.random.split(key)
        jax_setattr(self, "key", new_key)
        return returned_key


rng = StatefulRNG(jax.random.key(0))


def jitted():
    rng.split()
    rng.split()


jax.jit(jitted)()
```

(n.b. this is not code or a style I intend to use; I was just trying to experiment with `experimental.attrs` to wrap my head around it)

Output on my laptop (`JAX_TRACEBACK_FILTERING=off` [version attached](https://github.com/google/jax/files/14500467/attrs_bug.log.txt))

```
❯ JAX_PLATFORMS=cpu python attrs_bug.py
Traceback (most recent call last):
  File "/Users/nelhage/attrs_bug.py", line 27, in <module>
    jax.jit(jitted)()
  File "/Users/nelhage/attrs_bug.py", line 24, in jitted
    rng.split()
  File "/Users/nelhage/attrs_bug.py", line 14, in split
    new_key, returned_key = jax.random.split(key)
                            ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/jax/_src/random.py", line 295, in split
    return _return_prng_keys(wrapped, _split(typed_key, num))
                                      ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/jax/_src/random.py", line 281, in _split
    return prng.random_split(key, shape=shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/jax/_src/prng.py", line 631, in random_split
    return random_split_p.bind(keys, shape=shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: a jaxpr variable must be created only once per tracer
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```


The problem arises from these lines in `attrs.py`:

https://github.com/google/jax/blob/f9e20d58754283de87b2ed35cc9df58bcdff2073/jax/experimental/attrs.py#L70-L71

The `tracer` we create is referenced only by way of the `setattr` on the following line. If a future mutation to the tracked object drops that reference, the tracer will be freed, and its address (`id`) may be reused by a future object. If a new `Tracer` then happens to get allocated at the same address, the `tracer_to_var` map will still contain the old mapping and result in the observed problem.

We can hackily test that theory by making sure to manually retain all the tracers inside of our reproducer; this version works reliably for me:

```python
import jax

from jax.experimental.attrs import jax_getattr, jax_setattr


class StatefulRNG:
    key: jax.Array

    def __init__(self, key: jax.Array):
        self.key = key

    def split(self, l) -> jax.Array:
        key = jax_getattr(self, "key")
        l.append(key)
        new_key, returned_key = jax.random.split(key)
        l.append(new_key)
        l.append(returned_key)
        jax_setattr(self, "key", new_key)
        return returned_key


rng = StatefulRNG(jax.random.key(0))


def jitted():
    l = []
    rng.split(l)
    rng.split(l)
    l.clear()


jax.jit(jitted)()
```

### System info (python version, jaxlib version, accelerator, etc.)

```
❯ JAX_PLATFORMS='cpu' python -c 'import jax; jax.print_environment_info()'
jax:    0.4.25
jaxlib: 0.4.25
numpy:  1.24.4
python: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:37:07) [Clang 15.0.7 ]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Darwin', node='Nelson-Elhage-MacBook', release='23.3.0', version='Darwin Kernel Version 23.3.0: Wed Dec 20 21:30:27 PST 2023; root:xnu-10002.81.5~7/RELEASE_ARM64_T8103', machine='arm64')
```


Discussion:
Brilliant, thank you for the clear diagnosis!

I think we should probably persist all the Tracers we create during jaxpr tracing. That's usually what we do with [this attribute on the builder](https://github.com/google/jax/blob/67e3542d326d2bb6e4cc958e134ace4baaa152ec/jax/_src/interpreters/partial_eval.py#L1760), and we [append to it when we make new tracers in the usual path](https://github.com/google/jax/blob/67e3542d326d2bb6e4cc958e134ace4baaa152ec/jax/_src/interpreters/partial_eval.py#L2001), but as you point out we neglected to do it on the attrs path.

Resolving patch:
diff --git a/jax/experimental/attrs.py b/jax/experimental/attrs.py
--- a/jax/experimental/attrs.py
+++ b/jax/experimental/attrs.py
@@ -72,6 +72,7 @@ def _ensure_tracked(trace: pe.DynamicJaxprTrace, obj: Any, attr: str):
     frame.attrs_tracked.append((obj, attr))
     frame.attrs_inits.append(init_val)
     frame.attrs_vars.append(var)
+    frame.tracers.append(tracer)
 pe.DynamicJaxprTrace._ensure_tracked = _ensure_tracked
 
 def _getattr_staging(trace, *, obj, attr):


Test patch:
diff --git a/tests/attrs_test.py b/tests/attrs_test.py
--- a/tests/attrs_test.py
+++ b/tests/attrs_test.py
@@ -45,6 +45,7 @@ class Thing:
 
 class AttrsTest(jtu.JaxTestCase):
 
+
   @parameterized.parameters([True, False])
   def test_jit_basic(self, jit: bool):
     thing = Thing(1.0)
@@ -228,6 +229,28 @@ def f(obj, x):
     self.assertAllClose(thing.x, 3.0, check_dtypes=False)
     self.assertEqual(count, 1)
 
+  def test_tracer_lifetime_bug(self):
+    # regression test for https://github.com/google/jax/issues/20082
+    class StatefulRNG:
+      key: jax.Array
+
+      def __init__(self, key: jax.Array):
+        self.key = key
+
+      def split(self) -> jax.Array:
+        key = jax_getattr(self, "key")
+        new_key, returned_key = jax.random.split(key)
+        jax_setattr(self, "key", new_key)
+        return returned_key
+
+    rng = StatefulRNG(jax.random.key(0))
+
+    def jitted():
+      rng.split()
+      rng.split()
+
+    jax.jit(jitted)()  # don't crash
+
 
 class AttrsJVPTest(jtu.JaxTestCase):
 
@@ -533,5 +556,6 @@ def f_ref(x, y, z, w):
     self.assertAllClose(attr_cotangents[(thing2, 'x')], attr_cotangents_ref[1],
                         check_dtypes=False)
 
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())