Ionel Gog ec279f9c54 Add config option to log or fatal when jax.Arrays are GCed. ...

Introduces `jax.config.array_garbage_collection_guard`, which is a tristate config for setting up a `jax.Array` garbage collection guard. The possible configs are:
* allow: `jax.Array`s are allowed to be garbage collected. This is the default value.
* log: whenever a `jax.Array` is GCed a log entry is generated with the array's traceback.
* fatal: fatal crash when a `jax.Array` is GCed. This is meant to be used for mature code bases that do tight memory management, and are reference cycle free.

PiperOrigin-RevId: 687003464

2024-10-17 12:23:16 -07:00

..

mlir

Partially rolling forward #22998

2024-09-02 04:44:47 -07:00

__init__.py

Add config option to log or fatal when jax.Arrays are GCed.

2024-10-17 12:23:16 -07:00

BUILD

Clean up BUILD files.

2024-08-26 09:11:17 -07:00

mosaic_gpu.py

Make CuDNN SDPA API work with JAX with a CUDA plugin configuration.

2024-06-06 12:09:19 -04:00

triton.py

Do not import from lowering in tests/pallas/pallas_test.py

2024-05-10 14:25:17 -07:00