Changes:
- Fix unnecessary generator
- Iterate dictionary directly instead of calling .keys()
- Remove global statement at the module level
- Use list() instead of a list comprehension
- Use with statement to open the file
- Merge isinstance calls
Updates XLA to 00afc7bb81.
The new XLA release removes the use of protocol buffers from the XLA client. Fixes#349.
Add backward compatibility shims to jaxlib to allow older jax releases to still work on an up to date jaxlib.
The new XLA release also incorporates a fix that avoids a host-device copy for every iteration of a `lax.fori_loop()` on GPU. Fixes#402.
Add a new jaxlib.__version__ field, change jax/jaxlib compatibility logic to check for it.
