15 Pull requests merged by 4 people

• [nnx] improve randomness guide

  #4281 merged Oct 10, 2024

• Update to Flax NNX MNIST tutorial

  #4277 merged Oct 10, 2024

• Update Why Flax NNX guide

  #4262 merged Oct 10, 2024

• Update to Flax NNX Model Surgery

  #4276 merged Oct 10, 2024

• Update Flax NNX Landing Page

  #4274 merged Oct 9, 2024

• [nnx] add FSDP toy example with custom optimizer

  #4183 merged Oct 9, 2024

• Add flatten to nnx.graph autosummary in graph.rst

  #4255 merged Oct 8, 2024

• remove flax-nnx urls

  #4267 merged Oct 8, 2024

• [nnx] add custom_vjp to docs

  #4266 merged Oct 8, 2024

• [nnx] revive TrainState toy example

  #4226 merged Oct 8, 2024

• Fix typo: impost to import

  #4256 merged Oct 8, 2024

• [nnx] add cloudpickle support

  #4253 merged Oct 8, 2024

• Add why nnx

  #4240 merged Oct 8, 2024

• Update to Flax NNX Transforms doc

  #4264 merged Oct 8, 2024

• Add checkpointing guide to website index.

  #4263 merged Oct 7, 2024

16 Pull requests opened by 6 people

• [nnx] fix ToLinen kwargs

  #4270 opened Oct 8, 2024

• Porting RNN from Linen to NNX

  #4272 opened Oct 8, 2024

• support passing arguments directly to the struct.dataclass decorator

  #4275 opened Oct 9, 2024

• [nnx] add PrefixMapping

  #4278 opened Oct 9, 2024

• Update Flax NNX Randomness

  #4279 opened Oct 10, 2024

• Fixed nnx examples bad links in the README.md

  #4282 opened Oct 10, 2024

• Fix Philosophy links, update copyright year,

  #4283 opened Oct 10, 2024

• Update Flax NNX Glossary

  #4284 opened Oct 10, 2024

• Remove notebook exceptions in `docs_nnx` doctest

  #4285 opened Oct 10, 2024

• Update Flax NNX vs JAX Transformations guide

  #4286 opened Oct 10, 2024

• [nnx] state filters

  #4288 opened Oct 11, 2024

• Update Flax Evolution from Linen to NNX guide

  #4289 opened Oct 11, 2024

• [nnx] improve init_cache docs

  #4291 opened Oct 12, 2024

• update version to 0.10.0

  #4292 opened Oct 12, 2024

• [nnx] flaxlib flatten

  #4293 opened Oct 13, 2024

• Ugrade Flax NNX Checkpointing guide

  #4294 opened Oct 14, 2024

2 Issues closed by 2 people

• CustomVJP example with NNX

  #4265 closed Oct 8, 2024

• Make a guide for scan over layers

  #2752 closed Oct 8, 2024

2 Issues opened by 2 people

• NNX documentation missing pooling operations

  #4271 opened Oct 8, 2024

• Add nnx.shard_map

  #4261 opened Oct 7, 2024

7 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 RNN (NNX)

  #4259 commented on Oct 7, 2024 • 0 new comments

• SPMD for initializing model using nnx.jit

  #4129 commented on Oct 8, 2024 • 0 new comments

• AttributeError: module 'jax.config' has no attribute 'define_bool_state' when running `big_vision` on `tpu-vm-base`

  #3180 commented on Oct 9, 2024 • 0 new comments

• Upgrade Flax NNX Bridge API doc

  #4198 commented on Oct 8, 2024 • 0 new comments

• [nnx] fix mnist_tutorial link

  #4204 commented on Oct 10, 2024 • 0 new comments

• proxy more methods for `nnx.Variable`

  #4234 commented on Oct 8, 2024 • 0 new comments

• Stackless yashful

  #4246 commented on Oct 11, 2024 • 0 new comments