George Necula e643f98558 [shape_poly] Reimplement the shape constraint checking using shape assertions. ...

Most of the functionality is for the JAX native serialization case.
This relies on newly added functionality to xla_extension.refine_polymorphic_shapes
that handles custom calls @static_assertion.

As a beneficial side-effect now we get shape constraint checking for jax2tf
graph serialization when the resulting function is executed in graph mode.

2023-07-19 09:56:33 +03:00

..

__init__.py

Reapply: move jax.interpreters.ad to jax._src.interpreters.ad

2023-02-02 09:29:05 -08:00

ad.py

[shape_poly] Cleanups for the shape polymorphism APIs.

2023-06-30 15:56:57 +03:00

batching.py

Fix silly type error involving dims_out sometimes being a thunk and sometimes not.

2023-07-15 05:13:43 -07:00

mlir.py

[shape_poly] Reimplement the shape constraint checking using shape assertions.

2023-07-19 09:56:33 +03:00

partial_eval.py

[shape_poly] More cleanup for the internal APIs for shape polymorphism.

2023-07-13 16:37:53 +03:00

pxla.py

Enable passing fdo_profile in compiler_options in pxla.py

2023-07-18 14:18:28 -07:00

xla.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00