Peter Hawkins cb4abe754a [MHLO] Separate registrations for collective and initial_style primitives from the XLA translation rule registration. ...

Change in preparation for removing XLA translation rules for many primitives. However, even after the MHLO switch we still need to tag collective and initial_style primitives.

PiperOrigin-RevId: 441474701

2022-04-13 07:26:26 -07:00

..

__init__.py

[JAX] Move traceback_util.py into jax._src.

2020-11-04 09:02:59 -08:00

ad.py

prototyping dynamic shapes

2022-04-11 22:10:47 -07:00

batching.py

prototyping dynamic shapes

2022-04-11 22:10:47 -07:00

masking.py

Adds simple effect types to jaxprs

2022-04-11 11:50:41 -07:00

mlir.py

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -07:00

partial_eval.py

Add tests for higher order primitives

2022-04-12 18:12:44 -07:00

pxla.py

Fix the pytype error. PyType is looking for a __init__ method. This does not change the behavior of the class.

2022-04-12 09:36:28 -07:00

sharded_jit.py

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -07:00

xla.py

[MHLO] Separate registrations for collective and initial_style primitives from the XLA translation rule registration.

2022-04-13 07:26:26 -07:00