Jake VanderPlas ca39457ea9 JEX: move jax.linear_util to jax.extend.linear_util

2023-08-30 18:32:12 -07:00

..

mosaic

Load Python rules from rules_python

2023-08-24 10:22:57 -07:00

triton

Load Python rules from rules_python

2023-08-24 10:22:57 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

Load Python rules from rules_python

2023-08-24 10:22:57 -07:00

core.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

indexing.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

pallas_call.py

JEX: move jax.linear_util to jax.extend.linear_util

2023-08-30 18:32:12 -07:00

primitives.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

utils.py

[Pallas] Add support for lax.fori_loop to Mosaic lowering

2023-08-04 13:43:43 -07:00