Matthew Johnson d21b958f30 add some simple iree tests ...

This passes, though two of the interesting tests fail with what might be IREE
bugs (and so are currently skipped):

```shell
JAX_PLATFORMS='iree' pytest -n auto tests/core_test.py tests/api_test.py -k Dynamic
```

2022-04-14 10:55:00 -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

add some simple iree tests

2022-04-14 10:55:00 -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