6 Pull requests merged by 3 people

• Fix for removal of JAX DeviceLocalLayout.

  #127 merged Jun 23, 2025

• Change for DeviceLocalLayout in JAX 0.6.3.

  #126 merged Jun 23, 2025

• Update explicit layouts to use JAX 0.6.2 layout.Format.

  #125 merged Jun 20, 2025

• Update references of jax_layout.Layout to jax_layout.Format.

  #124 merged Jun 19, 2025

• Add FTRL and Adam optimizers.

  #123 merged Jun 17, 2025

• Checkpoints

  #122 merged Jun 17, 2025