Commits

Branch selector

 main

User selector

All users

All time

Commit History

Commits on Mar 13, 2026

• Copybara import of the project:

  Show description for 4183e15
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  4183e15

  Copy full SHA for 4183e15

• add tree-mode-nnx FLIP

  Show description for 74e9303
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  74e9303

  Copy full SHA for 74e9303

• add graph_updates to pmap, grad, and value_and_grad

  Show description for 90df65b
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  90df65b

  Copy full SHA for 90df65b

• add graph_updates to fori_loop

  Show description for a2aa0eb
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  a2aa0eb

  Copy full SHA for a2aa0eb

• Merge pull request #5326 from google:metrics-mask

  Show description for 4cf2b8c
  

  Flax Authors

  committed
  4cf2b8c

  Copy full SHA for 4cf2b8c

• add graph_updates to while_loop

  Show description for a8ecd06
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  a8ecd06

  Copy full SHA for a8ecd06

Commits on Mar 12, 2026

• add graph_updates for scan

  Show description for 56471a4
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  56471a4

  Copy full SHA for 56471a4

• Merge pull request #5270 from samanklesaria:split_method

  Show description for d4653e5
  

  Flax Authors

  committed
  d4653e5

  Copy full SHA for d4653e5

• add graph_updates to vmap

  Show description for 35972c2
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  35972c2

  Copy full SHA for 35972c2

• Remove jax/tools/colab_tpu.py.

  Show description for 8ed67b8
  

  hawkinsp

  authored and

  Flax Authors

  committed
  8ed67b8

  Copy full SHA for 8ed67b8

• Added support for data masking in Average, Accuracy and MultiMetric

  vfdev-5

  committed
  5074fb6

  Copy full SHA for 5074fb6

• add graph_updates argument in shard_map

  Show description for 115dcf5
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  115dcf5

  Copy full SHA for 115dcf5

• Add split method to RngStream

  samanklesaria

  committed
  696328a

  Copy full SHA for 696328a

• add graph_updates argument to jit

  Show description for a1899fd
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  a1899fd

  Copy full SHA for a1899fd

Commits on Mar 9, 2026

• Allow DenyList to be compared.

  Show description for 22468a1
  

  levskaya

  authored and

  Flax Authors

  committed
  22468a1

  Copy full SHA for 22468a1

• Merge pull request #5259 from ayulockedin:feat/gqa-support-mha

  Show description for 9b39916
  

  Flax Authors

  committed
  9b39916

  Copy full SHA for 9b39916

Commits on Mar 7, 2026

• support out_sharding mapping in set_metadata

  Show description for 4614ac2
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  4614ac2

  Copy full SHA for 4614ac2

Commits on Mar 6, 2026

• fix eval_shape's _to_variable

  Show description for 8335dfe
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  8335dfe

  Copy full SHA for 8335dfe

• Merge pull request #5247 from samanklesaria:eval_shape_sharding

  Show description for 49e34d1
  

  Flax Authors

  committed
  49e34d1

  Copy full SHA for 49e34d1

Commits on Mar 5, 2026

• Resolve code review comments

  samanklesaria

  committed
  4e222b6

  Copy full SHA for 4e222b6

• Support explicit sharding for nnx.eval_shape

  samanklesaria

  committed
  6bce143

  Copy full SHA for 6bce143

• Change 'sharding' name to 'out_sharding'

  samanklesaria

  committed
  3e3c459

  Copy full SHA for 3e3c459

• Added sharding propagation support in nnx.eval_shape

  

  vfdev-5

  authored and

  samanklesaria

  committed
  7035157

  Copy full SHA for 7035157

• avoid __eq__ in test_iter_graph

  Show description for 39b7613
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  39b7613

  Copy full SHA for 39b7613

• add recursive_map example

  Show description for b6a683e
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  b6a683e

  Copy full SHA for b6a683e

Commits on Mar 4, 2026

• Merge pull request #5237 from mohsinm-dev:fix/pytreenode-generic-parameters

  Show description for 2970c6f
  

  Flax Authors

  committed
  2970c6f

  Copy full SHA for 2970c6f

• Merge pull request #5226 from google:nnx-optimizer-removed-redund

  Show description for ef6e4b1
  

  Flax Authors

  committed
  ef6e4b1

  Copy full SHA for ef6e4b1

• Merge pull request #5292 from samanklesaria:no_optax_pin

  Show description for 81a7e42
  

  Flax Authors

  committed
  81a7e42

  Copy full SHA for 81a7e42

• add with_attributes

  Show description for bedfdae
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  bedfdae

  Copy full SHA for bedfdae

• Annotate `wrapped_fn`'s `self` argument in `decorator_lift_transform` for trace-ability

  Show description for 3229c96
  

  jimlinntu

  authored and

  Flax Authors

  committed
  3229c96

  Copy full SHA for 3229c96

• simplify set_view API

  Show description for 8d7f486
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  8d7f486

  Copy full SHA for 8d7f486

• simplify nnx.display for Variable

  Show description for bba63ce
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  bba63ce

  Copy full SHA for bba63ce

• update version to 0.12.5

  Show description for 6ec76a2
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  6ec76a2

  Copy full SHA for 6ec76a2

• remove experimental Variable.mutable property

  Show description for fc24589
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  fc24589

  Copy full SHA for fc24589

• improve error messages for tree mode errors

  Show description for 9c33569
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  9c33569

  Copy full SHA for 9c33569

Pagination

Previous

Next