8 Pull requests merged by 3 people

• deflake a test

  #4663 merged Mar 28, 2025

• Fix lazy_init typo

  #4657 merged Mar 27, 2025

• Test runner ubuntu upgrade 24.04

  #4659 merged Mar 27, 2025

• Fix flaky vmap test tolerance.

  #4653 merged Mar 26, 2025

• Use jax.nn.dot_product_attention when possible

  #4649 merged Mar 25, 2025

• Roll up package jax version and uv.lock

  #4648 merged Mar 25, 2025

• [nnx] support Array leaves in graph nodes

  #4612 merged Mar 24, 2025

• discord release webhook

  #4646 merged Mar 24, 2025

6 Pull requests opened by 3 people

• roll back #4646

  #4647 opened Mar 24, 2025

• Added identity export from JAX

  #4652 opened Mar 25, 2025

• add fast_flatten

  #4654 opened Mar 26, 2025

• v0.10.5

  #4656 opened Mar 26, 2025

• add experimental support for mutable array

  #4661 opened Mar 28, 2025

• [nnx] refactor remat

  #4662 opened Mar 28, 2025

2 Issues closed by 1 person

• Typo bug in nnx/bridge/wrappers.py

  #4650 closed Mar 27, 2025

• Several tests failing with `AttributeError: module 'jax.api_util' has no attribute 'debug_info'`

  #4585 closed Mar 25, 2025

5 Issues opened by 5 people

• Using flax.linen.intercept_methods on an nnx module

  #4666 opened Mar 29, 2025

• Adding `dot_general` to `Einsum` Module

  #4664 opened Mar 28, 2025

• Filtering with `nnx.PathContains('')` returns an empty state

  #4660 opened Mar 27, 2025

• Issue with NNX custom_vjp in activation quantization implementation

  #4651 opened Mar 25, 2025

• Distribute inference on a TPU

  #4644 opened Mar 24, 2025

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

• Issue with Optimizer Update in A2C Network with Optax Body:

  #4391 commented on Mar 26, 2025 • 0 new comments

• Sequential API is not that much powerful.

  #4642 commented on Mar 27, 2025 • 0 new comments

• Gradient Checkpointing causes model to compute junk results (NNX)

  #4626 commented on Mar 28, 2025 • 0 new comments

• Add some methods on Sequential container

  #4565 commented on Mar 25, 2025 • 0 new comments

• [nnx] Add NNX WeightNorm.

  #4568 commented on Mar 26, 2025 • 0 new comments

• [nnx] remove deprecated APIs

  #4627 commented on Mar 25, 2025 • 0 new comments

• [nnx] refactor GraphDef

  #4630 commented on Mar 25, 2025 • 0 new comments

• [nnx] flaxlib types

  #4639 commented on Mar 25, 2025 • 0 new comments

• Apply vmap jnp.pad over the batch in REFLECT convolution mode

  #4640 commented on Mar 26, 2025 • 0 new comments