Peter Hawkins f1ea67117e Split name_stack out of mlir.ModuleContext. ...

A unique name_stack is built for every equation, which means that we're constantly rebuilding ModuleContext objects, even though the lifetime of almost everything else (naturally) is the Module scope. Split name_stack into an object that is threaded separately, including as part of mlir.LoweringRuleContext.

PiperOrigin-RevId: 608594374

2024-02-20 07:17:23 -08:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

ad.py

Remove unnecessary jax.config imports

2024-02-19 06:21:33 -08:00

batching.py

del add_any_p and zeros_like_p, replace aval-dispatched traceable

2023-12-21 17:04:21 -08:00

mlir.py

Split name_stack out of mlir.ModuleContext.

2024-02-20 07:17:23 -08:00

partial_eval.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

pxla.py

Delete mesh_sharding_specs from JAX

2024-01-02 11:14:39 -08:00

xla.py

Partial rollback xla translation APIs that were removed in 0.4.24 release

2024-02-15 13:32:33 -08:00