5 Pull requests merged by 3 people

• Rollback of attention_bias parameter, because the change overrides the attention bias for injected attention functions.

  #4703 merged Apr 14, 2025

• Unit tests for attention_bias parameter to MultiHeadDotProductAttention. Add parameter to all overloads to make pytype happy.

  #4702 merged Apr 14, 2025

• Add attention_bias parameter to MultiHeadDotProductAttention.

  #4694 merged Apr 11, 2025

• [nnx] remove deprecated APIs

  #4627 merged Apr 10, 2025

• [NVIDIA] Support FP8 Einsum Op

  #4686 merged Apr 9, 2025

3 Pull requests opened by 2 people

• Internal change

  #4696 opened Apr 9, 2025

• [nnx] add JitWrapped API

  #4699 opened Apr 10, 2025

• Add `attention_bias` parameter to `MultiHeadDotProductAttention`.

  #4701 opened Apr 13, 2025

1 Issue closed by 1 person

• TypeError: '<' not supported between instances of 'VariableState' and 'VariableState'

  #4676 closed Apr 14, 2025

2 Issues opened by 2 people

• [NNX] Can't use while_loop with vmap where each loop count is different across vmapped dim.

  #4700 opened Apr 10, 2025

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

  #4695 opened Apr 9, 2025

8 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.

• Support for exporting models in the ONNX format

  #4430 commented on Apr 11, 2025 • 0 new comments

• Using flax.linen.intercept_methods on an nnx module

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

• Fixes #1925 non-str dict keys not suppoted in module state

  #4563 commented on Apr 9, 2025 • 0 new comments

• Implement spectral and instance norm in NNX

  #4623 commented on Apr 11, 2025 • 0 new comments

• [nnx] refactor GraphDef

  #4630 commented on Apr 14, 2025 • 0 new comments

• [nnx] flaxlib types

  #4639 commented on Apr 14, 2025 • 0 new comments

• add fast_flatten

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

• add experimental support for mutable array

  #4661 commented on Apr 9, 2025 • 0 new comments