Commits

Breadcrumbs

History for

onmain

User selector

All users

All time

Commit History

Commits on Apr 10, 2026

• Fix for change to `jax.Array.__contains__`

  Show description for 64aa65a
  

  Jake VanderPlas

  authored and

  Flax Authors

  committed
  64aa65a

  Copy full SHA for 64aa65a

Commits on Mar 2, 2026

• No public description

  Show description for 5486a63
  

  danielsuo

  authored and

  Flax Authors

  committed
  5486a63

  Copy full SHA for 5486a63

Commits on Jan 12, 2026

• Merge pull request #5145 from google:fix-more-typos

  Show description for 73e7985
  

  Flax Authors

  committed
  73e7985

  Copy full SHA for 73e7985

Commits on Jan 6, 2026

• [pmap] Avoid degraded performance under the new `jax.pmap`.

  Show description for 5d90851
  

  danielsuo

  authored and

  Flax Authors

  committed
  5d90851

  Copy full SHA for 5d90851

Commits on Jan 5, 2026

• Fixed typos in the docstrings using antigravity

  vfdev-5

  committed
  b006c4c

  Copy full SHA for b006c4c

Commits on Sep 30, 2025

• use jnp.stack instead of np.stack

  vfdev-5

  authored
  c08e60b

  Copy full SHA for c08e60b

Commits on Mar 12, 2025

• Fixed typo in documentation for `save_checkpoint_multiprocess`.

  Show description for fefc4c9
  

  agutkin

  authored and

  Flax Authors

  committed
  fefc4c9

  Copy full SHA for fefc4c9

Commits on Jan 23, 2025

• Merge pull request #4499 from google:nnx-improve-module-docs

  Show description for e3bcc44
  

  Flax Authors

  committed
  e3bcc44

  Copy full SHA for e3bcc44

• [nnx] improve Module docs

  cgarciae

  committed
  bdcc33a

  Copy full SHA for bdcc33a

• Remove usages of orbax_utils.save_args_from_target, as this function does nothing (it used to control a checkpointing behavior that has since been optimized away).

  Show description for b5d4ed8
  

  cpgaffney1

  authored and

  Flax Authors

  committed
  b5d4ed8

  Copy full SHA for b5d4ed8

Commits on Oct 1, 2024

• Internal change

  Show description for 0b34068
  

  IvyZX

  authored and

  Flax Authors

  committed
  0b34068

  Copy full SHA for 0b34068

Commits on Aug 12, 2024

• Add logging to track deprecated codepaths.

  Show description for e600351
  

  cpgaffney1

  authored and

  Flax Authors

  committed
  e600351

  Copy full SHA for e600351

Commits on Jul 23, 2024

• fix restore_checkpoint example in docstring

  Show description for 46dc51e
  

  Cristian Garcia

  authored and

  Flax Authors

  committed
  46dc51e

  Copy full SHA for 46dc51e

Commits on Jul 18, 2024

• Ignore Orbax warning in deprecated flax.training and unbreak head

  IvyZX

  committed
  b0ebefc

  Copy full SHA for b0ebefc

Commits on Jul 11, 2024

• Remove some outdated backward-compatibility code.

  Show description for 5c97143
  

  IvyZX

  authored and

  Flax Authors

  committed
  5c97143

  Copy full SHA for 5c97143

Commits on Jul 10, 2024

• Modify Flax checkpointing in preparation for cl/650338576.

  Show description for e59b14e
  

  cpgaffney1

  authored and

  Flax Authors

  committed
  e59b14e

  Copy full SHA for e59b14e

Commits on Jul 8, 2024

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

  Show description for 96cf649
  

  dubey

  authored and

  Flax Authors

  committed
  96cf649

  Copy full SHA for 96cf649

Commits on Jul 5, 2024

• [nnx] stabilize unsafe_pytree

  cgarciae

  committed
  4ed137c

  Copy full SHA for 4ed137c

Commits on Jul 3, 2024

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

  Show description for 0fb1777
  

  cpgaffney1

  authored and

  Flax Authors

  committed
  0fb1777

  Copy full SHA for 0fb1777

Commits on Jun 27, 2024

• upgrade python to 3.10 + use pyupgrade

  cgarciae

  committed
  97ca05d

  Copy full SHA for 97ca05d

Commits on Apr 23, 2024

• Remove `ocdbt_merge` and `restore_with_serialized_types` options from `PyTreeCheckpointHandler` as they are not in use anywhere.

  Show description for fa119fd
  

  cpgaffney1

  authored and

  Flax Authors

  committed
  fa119fd

  Copy full SHA for fa119fd

Commits on Mar 15, 2024

• Make `TrainState`'s `step` possibly jax.Array. This makes `replicate` valid for type checking.

  Show description for 718aa8c
  

  IvyZX

  authored and

  Flax Authors

  committed
  718aa8c

  Copy full SHA for 718aa8c

Commits on Feb 26, 2024

• Use `jax.tree_util.tree_map` instead of deprecated `jax.tree_map`.

  Show description for 0797d05
  

  IvyZX

  authored and

  Flax Authors

  committed
  0797d05

  Copy full SHA for 0797d05

Commits on Jan 30, 2024

• added flax.typing

  chiamp

  committed
  49ba155

  Copy full SHA for 49ba155

Commits on Jan 22, 2024

• Merge pull request #3621 from google:nontransparent

  Show description for 7e9bb76
  

  Flax Authors

  committed
  7e9bb76

  Copy full SHA for 7e9bb76

Commits on Dec 9, 2023

• fixed double backquote code font

  chiamp

  committed
  b1f6aa2

  Copy full SHA for b1f6aa2

Commits on Nov 10, 2023

• added example docstrings

  chiamp

  committed
  72e217e

  Copy full SHA for 72e217e

Commits on Nov 2, 2023

• restructured doc folders

  chiamp

  committed
  1383ea3

  Copy full SHA for 1383ea3

• added pre-commit hook that sort imports and formats

  chiamp

  committed
  6fa3839

  Copy full SHA for 6fa3839

Commits on Oct 27, 2023

• added has_improved field to EarlyStopping

  chiamp

  committed
  b620117

  Copy full SHA for b620117

Commits on Oct 26, 2023

• Merge pull request #3435 from kaixih:update_fp8_support

  Show description for 738078c
  

  Flax Authors

  committed
  738078c

  Copy full SHA for 738078c

• Updated docstring for `axis_name` arg to say it's only used for `pmap`

  Show description for 5557649
  

  chiamp

  authored and

  Flax Authors

  committed
  5557649

  Copy full SHA for 5557649

Commits on Oct 24, 2023

• Update the fp8 support

  kaixih

  committed
  d24db61

  Copy full SHA for d24db61

Commits on Oct 17, 2023

• Add instrumentation for Flax checkpoint save and restore calls.

  Show description for 31eedf1
  

  Flax Team

  committed
  31eedf1

  Copy full SHA for 31eedf1

Commits on Oct 13, 2023

• Remove import checks that are no longer valid.

  Show description for 3db72f5
  

  IvyZX

  authored and

  Flax Authors

  committed
  3db72f5

  Copy full SHA for 3db72f5