jax-ml/jax jax-v0.9.1

JAX v0.9.1

on GitHub

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.