jax authors 82504194f7 Merge pull request #20718 from kkiningh:patch-3 ...

PiperOrigin-RevId: 624336436

2024-04-12 17:22:24 -07:00

..

mosaic

[Pallas TPU] Raise clearer NotImplementedError on vector -> scalar reductions.

2024-04-12 15:00:07 -07:00

triton

_cast() now takes JAX dtypes

2024-04-10 17:36:39 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Fix some load orderings for buildifier

2024-03-27 10:28:57 -07:00

core.py

Add dynamic grid support to emit_pipeline

2024-04-09 23:55:51 -07:00

pallas_call.py

pallas_call now has only one way to pass compiler_params=

2024-04-09 14:23:20 -07:00

primitives.py

Add type information to Pallas primatives.

2024-04-11 13:08:02 -07:00

utils.py

[Pallas TPU] Convert pattern_match_while_to_fori_loop to return (Jaxpr, str) rather than throw exceptions.

2024-04-09 14:04:25 -07:00