7 Pull requests merged by 4 people

• Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file.

  #4039 merged Jul 3, 2024

• [nnx] fix Variable overloads and add shape/dtype properties

  #4049 merged Jul 2, 2024

• Fix typo in nnx_basics doc

  #4047 merged Jul 2, 2024

• Add functool.wraps() annotation to flax.nn.jit.

  #4051 merged Jul 2, 2024

• Added support of NANOO fp8.

  #3993 merged Jul 2, 2024

• [linen] allows multiple compact methods

  #3808 merged Jul 1, 2024

• [nnx] LoRAParam inherits from Param

  #3988 merged Jul 1, 2024

6 Pull requests opened by 4 people

• Update nnx_basics.md

  #4050 opened Jul 2, 2024

• Unsets `profile_duration_ms` in Flax ResNet example to prioritize `num_profile_steps`.

  #4052 opened Jul 2, 2024

• Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file.

  #4055 opened Jul 4, 2024

• Allow metadata pass-through in flax.struct.field

  #4056 opened Jul 4, 2024

• [linen] generalize transform caching

  #4057 opened Jul 4, 2024

• [nnx] improve filters guide

  #4059 opened Jul 5, 2024

4 Issues closed by 3 people

• Opaque XLA crash when initializing model

  #4054 closed Jul 4, 2024

• New python venv fails to pip install the flax mnist example requirements

  #3990 closed Jul 3, 2024

• typo in nnx_basics.md

  #4046 closed Jul 2, 2024

• Documentation links 404

  #3917 closed Jul 1, 2024

2 Issues opened by 2 people

• Is there anyway to analyze activations in flax?

  #4058 opened Jul 5, 2024

• Significant performance difference of NNX relative to equinox

  #4045 opened Jun 30, 2024

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.

• [NVIDIA] Rename fp8 custom dtype to `fp32_max_grad`

  #3984 commented on Jul 2, 2024 • 3 new comments

• [nnx] stabilize unsafe_pytree

  #4030 commented on Jul 5, 2024 • 3 new comments

• lstm error

  #4032 commented on Jul 2, 2024 • 1 new comment

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

  #3180 commented on Jul 3, 2024 • 1 new comment

• Update orbax_upgrade_guide.rst for async checkpointing usage examples

  #4036 commented on Jul 3, 2024 • 1 new comment

• Using `vmap` on the loss function change results

  #3973 commented on Jul 3, 2024 • 0 new comments

• [nnx] experimental transforms

  #3963 commented on Jul 4, 2024 • 0 new comments

• [linen] allow checkpoint to cache

  #3965 commented on Jul 4, 2024 • 0 new comments