jax-ml/jax jaxlib-v0.4.29

Jaxlib release v0.4.29

on GitHub

Bug fixes

• Fixed a bug where XLA sharded some concatenation operations incorrectly,
  which manifested as an incorrect output for cumulative reductions (#21403).
• Fixed a bug where XLA:CPU miscompiled certain matmul fusions
  (openxla/xla#13301).
• Fixes a compiler crash on GPU (#21396).

Deprecations

• jax.tree.map(f, None, non-None) now emits a DeprecationWarning, and will
  raise an error in a future version of jax. None is only a tree-prefix of
  itself. To preserve the current behavior, you can ask jax.tree.map to
  treat None as a leaf value by writing:
  jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None).