Jax: jax-v0.9.1 Release

Changes:

• JAX tracers that are not of Array type (e.g., of Ref type) will no longer report themselves to be instances of Array.
• Using jax.shard_map in Explicit mode will raise an error if the PartitionSpec of input does not match the PartitionSpec specified in in_specs. In other words, it will act like an assert instead of an implicit reshard. in_specs is an optional argument so you can omit specifying it and shard_map will infer the PartitionSpec from the argument. If you want to reshard your inputs, you can use jax.reshard on the arguments and then pass those args to shard_map.

New features:

• Added a debug config jax_compilation_cache_check_contents. If set, we miss when get() is called on a value that has not been put() by the current process, even if the value is actually in the disk cache. When a value is put(), we verify that its contents match.