Jake VanderPlas 40367a9eaf Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08:00

..

__init__.py

Removed the double re-exporting of Pallas GPU/TPU APIs

2024-06-10 05:59:09 -07:00

BUILD

Use JAX's default device instead of jax.devices()[0], if set.

2024-12-03 16:53:58 -08:00

core.py

Reverts f281c6f46475270a57a02416469226315377592c

2024-11-05 07:17:14 -08:00

error_handling.py

Use the current minimum jaxlib version for type checking on the CI

2024-10-10 12:46:15 +01:00

lowering.py

[Pallas TPU] Add WeirdOp to TPU dialect and add lowering for lax.is_finite

2024-12-10 16:38:04 -08:00

pallas_call_registration.py

jax.core: deprecate a number of APIs

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

pipeline.py

Use JAX's default device instead of jax.devices()[0], if set.

2024-12-03 16:53:58 -08:00

primitives.py

Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08:00

random.py

[Pallas] Add lowering for threefry PRNG.

2024-10-09 14:48:26 -07:00

verification.py

[Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction

2024-12-05 07:06:19 -08:00