3 Pull requests merged by 3 people

• Make fully replicated array before saving checkpoints for examples that use pmap.

  #4707 merged Apr 16, 2025

• [nnx] refactor GraphDef

  #4630 merged Apr 15, 2025

• Add custom einsum op to Einsum()

  #4705 merged Apr 15, 2025

3 Pull requests opened by 2 people

• remove "nnx" collection in ToLinen

  #4708 opened Apr 16, 2025

• convert DenseGeneral to NNX

  #4710 opened Apr 17, 2025

• Roll up ml_collections version to fix head

  #4711 opened Apr 17, 2025

2 Issues closed by 2 people

• Adding `dot_general` to `Einsum` Module

  #4664 closed Apr 15, 2025

• `nnx.Optimizer.update` causes tracer leaking when used in a jitted function

  #4695 closed Apr 15, 2025

2 Issues opened by 2 people

• TraceContextError when JIT-compiling nnx.RNN with nnx.LSTMCell

  #4713 opened Apr 18, 2025

• [NNX] Can't use fori_loop inside cond with vmap.

  #4709 opened Apr 17, 2025

5 Unresolved conversations

Sometimes conversations happen on old items that aren’t yet closed. Here is a list of all the Issues and Pull Requests with unresolved conversations.

• [nnx] add JitWrapped API

  #4699 commented on Apr 18, 2025 • 3 new comments

• [nnx] flaxlib types

  #4639 commented on Apr 17, 2025 • 2 new comments

• Using flax.linen.intercept_methods on an nnx module

  #4666 commented on Apr 15, 2025 • 0 new comments

• Failed build: CI - with JAX nightly

  #4488 commented on Apr 18, 2025 • 0 new comments

• add fast_flatten

  #4654 commented on Apr 17, 2025 • 0 new comments