Matthew Johnson 4354f355a8 prototyping dynamic shapes ...

Co-authored-by: Dougal Maclaurin <dougalm@google.com>

2022-04-11 22:10:47 -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

prototyping dynamic shapes

2022-04-11 22:10:47 -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

Adds simple effect types to jaxprs

2022-04-11 11:50:41 -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