Concatenate 2 jax trace arrays

[BrandonGel]

Hey,

So I am trying to concatenate 2 trace arrays to make 1 new array, but I am getting this error. Not sure how to resolve it.

`action
Traced<ShapedArray(float32[20,256,12])>with<DynamicJaxprTrace(level=3/1)>`
`observation
Traced<ShapedArray(float32[20,256,465])>with<DynamicJaxprTrace(level=3/1)>`
`a = jnp.concat((observation,action),axis=-1)
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1906, in concat
    return jax.numpy.concatenate(arrays, axis=axis)
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1898, in concatenate
    arrays_out = [lax.concatenate(arrays_out[i:i+k], axis)
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1898, in <listcomp>
    arrays_out = [lax.concatenate(arrays_out[i:i+k], axis)
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 622, in concatenate
    return concatenate_p.bind(*operands, dimension=dimension)
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/core.py", line 422, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/core.py", line 1337, in find_top_trace
    top_tracer._assert_live()
  File "/home/bho36/envs/fill/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 1720, in _assert_live
    raise core.escaped_tracer_error(self, None)
jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float32[20,256,465] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
The function being traced when the value leaked was minibatch_step at /home/bho36/envs/fill/lib/python3.10/site-packages/fill/algorithms/gail/train.py:244 traced for scan.
------------------------------
The leaked intermediate value was created on line /home/bho36/FILL/scripts/train_gail_barkour_obstacle.py:78 (main). 
------------------------------
When the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
------------------------------
/home/bho36/.vscode-server/extensions/ms-python.debugpy-2024.0.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py:321 (run_path)
/home/bho36/.vscode-server/extensions/ms-python.debugpy-2024.0.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py:135 (_run_module_code)
/home/bho36/.vscode-server/extensions/ms-python.debugpy-2024.0.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py:124 (_run_code)
/home/bho36/FILL/scripts/train_gail_barkour_obstacle.py:106 (<module>)
/home/bho36/FILL/scripts/train_gail_barkour_obstacle.py:78 (main)
------------------------------

To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError`

[jakevdp]

This looks unrelated to concatenation, rather it seems your function is leaking a tracer, which generally means you're attempting to modify some global state within JIT-compiled code.

It will be difficult to help any further without more information; could you try to put together a minimal reproducible example that reproduces the error you're seeing?

[BrandonGel]

Ok, I found out my error. When using a python debugger with JAX, it would cause a trace memory leakage. So not using debugger, the code works fine.