Peter Hawkins f4adcc650f Set __slots__ on core.Trace subclasses. ...

This is easy to do and makes field accesses on Trace classes slightly faster.

2025-01-22 16:17:54 -05:00

..

__init__.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

_base.py

Replaced most usages of abc.ABC with util.StrictABC

2024-01-29 12:40:43 +00:00

_lowerings.py

jax.core: deprecate a number of APIs

2024-12-10 11:11:32 -08:00

ad.py

Re-pack gradients of jax.experimental.sparse.grad() to match original pytrees & test cases

2024-07-29 13:04:05 -04:00

api.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

bcoo.py

Rename out_type -> out_sharding parameter on einsum

2025-01-16 18:16:52 -08:00

bcsr.py

Rename out_type -> out_sharding parameter on einsum

2025-01-16 18:16:52 -08:00

coo.py

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00

csr.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

linalg.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

nm.py

jax.core: deprecate a number of APIs

2024-12-10 11:11:32 -08:00

random.py

update experimental.sparse to new-style typed keys

2024-03-07 12:40:10 -08:00

test_util.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

transform.py

Set __slots__ on core.Trace subclasses.

2025-01-22 16:17:54 -05:00

util.py

Rename out_type -> out_sharding parameter on einsum

2025-01-16 18:16:52 -08:00