Commits

Branch selector

 main

User selector

All users

All time

Commit History

Commits on Mar 27, 2026

• Replaced custom PriorityStr with standard tuple sorting in _graph_node_flatten. Python 3.13.11 optimized str comparison bypasses custom subclasses (causing TypeError), and used getattr in _setattr …

  Show description for c8e3d84
  

  Flax Team

  committed
  c8e3d84

  Copy full SHA for c8e3d84

Commits on Mar 26, 2026

• Add nnx.make_jaxpr.

  Show description for d539ce7
  

  Flax Team

  committed
  d539ce7

  Copy full SHA for d539ce7

Commits on Mar 24, 2026

• Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

  Show description for 35d911a
  

  Flax Authors

  committed
  35d911a

  Copy full SHA for 35d911a

• Merge pull request #5231 from samanklesaria:opt_sharding

  Show description for bb50f21
  

  Flax Authors

  committed
  bb50f21

  Copy full SHA for bb50f21

Commits on Mar 21, 2026

• fix transform_metadata

  Show description for a138d9f
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  a138d9f

  Copy full SHA for a138d9f

Commits on Mar 20, 2026

• add nnx.map and nnx.abstract_with_sharding

  Show description for 1572c84
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  1572c84

  Copy full SHA for 1572c84

• Merge pull request #5246 from samanklesaria:generalize_out_sharding

  Show description for 3bcc4a9
  

  Flax Authors

  committed
  3bcc4a9

  Copy full SHA for 3bcc4a9

• improve aliasing error msg

  Show description for 1a56700
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  1a56700

  Copy full SHA for 1a56700

• add graph node in prefix checks

  Show description for 9eb565c
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  9eb565c

  Copy full SHA for 9eb565c

• improve nnx.Dict error handling

  Show description for ebc8185
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  ebc8185

  Copy full SHA for ebc8185

• simplify SimpleScan

  Show description for 2965f95
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  2965f95

  Copy full SHA for 2965f95

• Do a few more cleanups after the pmap shmap merge deletion.

  Show description for f521e95
  

  yashk2810

  authored and

  Flax Authors

  committed
  f521e95

  Copy full SHA for f521e95

• Do a few more cleanups after pmap_shmap merge

  Show description for 50433ef
  

  hawkinsp

  authored and

  Flax Authors

  committed
  50433ef

  Copy full SHA for 50433ef

• update to version 0.12.6

  Show description for ee6d3c5
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  ee6d3c5

  Copy full SHA for ee6d3c5

• Do a few more cleanups after pmap_shmap merge

  Show description for b96a07a
  

  yashk2810

  authored and

  Flax Authors

  committed
  b96a07a

  Copy full SHA for b96a07a

Commits on Mar 19, 2026

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

  Show description for 2b3ffaf
  

  Flax Authors

  committed
  2b3ffaf

  Copy full SHA for 2b3ffaf

• Use with_sharding_constraint for Format application

  samanklesaria

  committed
  1c2cec4

  Copy full SHA for 1c2cec4

• Generalize out_sharding to work with NamedSharding and Format

  samanklesaria

  committed
  0399b9d

  Copy full SHA for 0399b9d

• fix jit_partial lower and compile

  Show description for 15683fd
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  15683fd

  Copy full SHA for 15683fd

• Merge pull request #5257 from samanklesaria:capture_intms

  Show description for e5a6ff6
  

  Flax Authors

  committed
  e5a6ff6

  Copy full SHA for e5a6ff6

Commits on Mar 18, 2026

• Move orbax-export to optional deps

  samanklesaria

  committed
  c4d1566

  Copy full SHA for c4d1566

• check aliases on all transform args and simplify apply_variable_updates

  Show description for 1fcc2a5
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  1fcc2a5

  Copy full SHA for 1fcc2a5

• remove aliases from nnx.graph in favor of nnx.compat

  Show description for b313d06
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  b313d06

  Copy full SHA for b313d06

• Introduce `manual_type: ManualAxisType` parameter on ShapedArray to track varying/unreduced/reduced and remove `vma: frozenset` parameter.

  Show description for 4ba66fc
  

  yashk2810

  authored and

  Flax Authors

  committed
  4ba66fc

  Copy full SHA for 4ba66fc

• improve error messages for tree mode duplicates check

  Show description for d930ccd
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  d930ccd

  Copy full SHA for d930ccd

• clean up custom_vjp's graph_updates=False section

  Show description for 93c430d
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  93c430d

  Copy full SHA for 93c430d

• add transform_metadata transform

  Show description for 1779706
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  1779706

  Copy full SHA for 1779706

Commits on Mar 17, 2026

• allow in/out_axes when graph_updates=False

  Show description for 4f58b35
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  4f58b35

  Copy full SHA for 4f58b35

• Add nnx.capture

  samanklesaria

  committed
  0e92624

  Copy full SHA for 0e92624

• support hijax and ref Variables in simple transforms

  Show description for a9059fc
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  a9059fc

  Copy full SHA for a9059fc

• don't allow Variable mutation in custom_vjp on differentiable arguments

  Show description for d812e60
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  d812e60

  Copy full SHA for d812e60

• add more tests that check for consistent aliasing in transforms

  Show description for df5c180
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  df5c180

  Copy full SHA for df5c180

• add compat module

  Show description for 5a79a96
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  5a79a96

  Copy full SHA for 5a79a96

Commits on Mar 16, 2026

• add graph_updates to eval_shape and checkify

  Show description for 6b358ed
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  6b358ed

  Copy full SHA for 6b358ed

• add graph_updates to vjp and jvp

  Show description for 4386571
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  4386571

  Copy full SHA for 4386571

Pagination

Previous

Next