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

register the lax module for stack trace filtering

2021-02-22 21:11:41 -08:00

ann.py

[JAX] Fix batch logic for approx_min/max_k

2022-04-08 13:50:36 -07:00

control_flow.py

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

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

convolution.py

[MLIR] Change signature of lowering rules.

2021-12-15 19:06:58 -08:00

fft.py

Support string fft_type values in lax.fft.

2022-02-03 08:52:38 -05:00

lax.py

prototyping dynamic shapes

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

linalg.py

Consolidation of hipsolver/cusolver APIs.

2022-04-07 01:46:43 +00:00

other.py

Use real Precision type for lax.PrecisionType

2022-03-04 04:21:25 -08:00

parallel.py

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

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

polar.py

Document jax.lax.Precision

2021-11-08 14:15:31 -08:00

qdwh.py

[linalg] Fix type promotion in QDWH.

2022-03-03 20:59:28 -08:00

slicing.py

prototyping dynamic shapes

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

svd.py

[JAX] Added jit-able singular value decomposition.

2022-02-03 11:16:55 -08:00

utils.py

prototyping dynamic shapes

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

windowed_reductions.py

fixup from 5415306: remove extraneous lines

2022-03-11 15:19:10 -08:00