6 Commits

Author SHA1 Message Date
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
jax authors
d5a1c64d13 Merge pull request #9623 from rsuderman:FixIreeJax
PiperOrigin-RevId: 430834942
2022-02-24 17:51:09 -08:00
Rob Suderman
23d5eb0896 IREE's get_default_device_assignment should return List[Device]
Previously return List[List[Device]] which is not how the function is used.
Updated to use the alternative overload.
2022-02-18 23:02:04 -08:00
Rob Suderman
50f4b5806e Fix Iree backend to for copy_to_device and executable results
Executable results can be a tuple, if so iterate over the entires.

Copy to device should just return the the IREE buffer as device buffer
management is still in progress.
2022-02-17 14:31:23 -08:00
Peter Hawkins
04d8b3544e [JAX:IREE] Fix bitrot in iree plugin for JAX.
The compiler driver appears not to be a thing in the IREE API any more. But it also appears not to be needed.

PiperOrigin-RevId: 424190139
2022-01-25 15:08:49 -08:00
Peter Hawkins
70b8a6a806 Add a prototype IREE backend for JAX.
This is to support experimentation with the combination of JAX/IREE. Many things do not work yet.

PiperOrigin-RevId: 409980064
2021-11-15 07:57:04 -08:00