The Wayback Machine - https://web.archive.org/web/20230215045407/https://github.com/google/flax/commits/main
Skip to content
Permalink
main
Switch branches/tags

Name already in use

A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?

Commits on Feb 15, 2023

  1. Flax layers are parametrizable with custom conv_general_dilated. 

    PiperOrigin-RevId: 509679602
    Flax Team committed Feb 15, 2023
    92b0cdd

Commits on Feb 13, 2023

  1. Remove fallbacks for old versions of JAX from Flax. 

    Flax's setup.py requires JAX >= 0.4.2, so these code paths are dead.

PiperOrigin-RevId: 509309356
    @hawkinsp
    hawkinsp authored and Flax Authors committed Feb 13, 2023
    323d6b8

Commits on Feb 9, 2023

  1. Flax layers are parametrizable with custom dot_general. 

    PiperOrigin-RevId: 508438464
    Flax Team committed Feb 9, 2023
    1f6b094

  2. Replace uses of deprecated JAX sharding APIs with their new names in … 

    …jax.sharding.

This change updates:
* {jax.experimental.maps.Mesh, jax.interpreters.pxla.Mesh} to jax.sharding.Mesh
* {jax.experimental.PartitionSpec, jax.experimental.pjit.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.pxla.PartitionSpec} to jax.sharding.PartitionSpec
* jax.experimental.maps.NamedSharding to jax.sharding.NamedSharding.

PiperOrigin-RevId: 508404789
    @hawkinsp
    hawkinsp authored and Flax Authors committed Feb 9, 2023
    26d46e4

  3. Merge pull request #2856 from levskaya:point_release4 

    PiperOrigin-RevId: 508228804
    Flax Authors committed Feb 9, 2023
    524629e

  4. Flax 0.6.5 point release

    @levskaya
    levskaya committed Feb 9, 2023
    5c6f426

Commits on Feb 8, 2023

  1. Merge pull request #2837 from cgarciae:improve-remat-docs 

    PiperOrigin-RevId: 508161859
    Flax Authors committed Feb 8, 2023
    6ea0aca

Commits on Feb 7, 2023

  1. Merge pull request #2816 from jheek:lazy-init 

    PiperOrigin-RevId: 507650305
    Flax Authors committed Feb 7, 2023
    24224bf

  2. Merge pull request #2811 from arcAman07:main 

    PiperOrigin-RevId: 507648060
    Flax Authors committed Feb 7, 2023
    c3c4a87

Commits on Feb 3, 2023

  1. Replace uses of deprecated JAX sharding APIs with their new names in … 

    …jax.sharding.

This change updates:
* {jax.experimental.maps.Mesh, jax.interpreters.pxla.Mesh} to jax.sharding.Mesh
* {jax.experimental.PartitionSpec, jax.experimental.pjit.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.pxla.PartitionSpec} to jax.sharding.PartitionSpec
* jax.experimental.maps.NamedSharding to jax.sharding.NamedSharding.

PiperOrigin-RevId: 506995236
    @hawkinsp
    hawkinsp authored and Flax Authors committed Feb 3, 2023
    398e170

  2. improve nn.checkpoint docs

    @cgarciae
    cgarciae committed Feb 3, 2023
    211c48e

  3. Merge pull request #2815 from chiamp:initializers_refactor 

    PiperOrigin-RevId: 506842072
    Flax Authors committed Feb 3, 2023
    06529c9

  4. Merge pull request #2827 from google:update-jax-dep 

    PiperOrigin-RevId: 506767322
    Flax Authors committed Feb 3, 2023
    8504e5d

  5. replaced zeros and ones initializers with zeros_init and ones_init bu… 

    …ilder functions
    @chiamp
    chiamp committed Feb 3, 2023
    2d32f28

Commits on Feb 2, 2023

  1. Merge pull request #2843 from IvyZX:meta 

    PiperOrigin-RevId: 506736298
    Flax Authors committed Feb 2, 2023
    c06cba2

  2. Make a separate Python module for Orbax-related utils. 

    PiperOrigin-RevId: 506731539
    @IvyZX
    IvyZX authored and Flax Authors committed Feb 2, 2023
    694560b

  3. Replace uses of deprecated JAX sharding APIs with their new names in … 

    …jax.sharding.

This change updates:
* {jax.experimental.maps.Mesh, jax.interpreters.pxla.Mesh} to jax.sharding.Mesh
* {jax.experimental.PartitionSpec, jax.experimental.pjit.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.pxla.PartitionSpec} to jax.sharding.PartitionSpec
* jax.experimental.maps.NamedSharding to jax.sharding.NamedSharding.

PiperOrigin-RevId: 506722866
    @hawkinsp
    hawkinsp authored and Flax Authors committed Feb 2, 2023
    431787b

  4. Re-box axis metadata in DenseGeneral's init functions

    @IvyZX
    IvyZX committed Feb 2, 2023
    ba87b46

  5. Apply suggestions from code review 

    Co-authored-by: Marc van Zee <marcvanzee@gmail.com>
    @jheek @marcvanzee
    jheek and marcvanzee committed Feb 2, 2023
    6ffdd83

  6. Merge pull request #2830 from chiamp:import_initializers 

    PiperOrigin-RevId: 506669714
    Flax Authors committed Feb 2, 2023
    aaba599

  7. Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code. 

    Uses jax.local_devices() as argument to jax.device_put_replicated and jax.device_put_sharded, removing the custom logic for older JAX versions.

PiperOrigin-RevId: 506667707
    @marcvanzee
    marcvanzee authored and Flax Authors committed Feb 2, 2023
    83a5b3b

  8. Clarify we are using the ReZero trick in our ResNet example. 

    PiperOrigin-RevId: 506666617
    @marcvanzee
    marcvanzee authored and Flax Authors committed Feb 2, 2023
    1b5b504

  9. Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code. 

    Uses jax.local_devices() as argument to jax.device_put_replicated and jax.device_put_sharded, removing the custom logic for older JAX versions.

PiperOrigin-RevId: 506573580
    @marcvanzee
    marcvanzee authored and Flax Authors committed Feb 2, 2023
    06323f2

  10. Apply suggestions from code review 

    Co-authored-by: Marc van Zee <marcvanzee@gmail.com>
    @jheek @marcvanzee
    jheek and marcvanzee committed Feb 2, 2023
    f4da05b

  11. Merge pull request #2833 from levskaya:relaxed_naming 

    PiperOrigin-RevId: 506463179
    Flax Authors committed Feb 2, 2023
    dfb55c4

  12. Introduce relaxed naming policy behind a feature flag. 

    Historically Flax had strong opinions about how variables and submodules
could be named, in the spirit of maintaining a single coherent namespace
amongst python attributes, variable, and submodule names.

It turns out we never really enforced that "single namespace" concept
correctly, and furthermore, that the motivations for such a thing were
very weak at best and any possible advantages left unrealized.

In practice, these name constraints are just really annoying to our users
and mostly just get in the way, especially with advanced configuration
systems.

This PR removes the pointless constraints:

 - "Adopted" modules brought in from the outside, as commonly occurs
   with configuration systems like fiddle, gin, etc. are now able to
   carry "scope names" which are not reset on adoption during module
   construction.
 - Variables in different collections in the same scope/module can have
   the same "scope name" in the variable nested dict.  Conflicts
   between submodule names and variable names are still enforced.
 - We no longer force variable "scope names" not to clash with
   python attributes on the module.  We only enforce non-collision
   within the backing scope itself.

Since some of these changes could potentially change checkpoints in
certain situations, this new behavior is released behind a flag
"flax_relaxed_naming" for the time being.
    @levskaya
    levskaya committed Feb 2, 2023
    54e342b

Commits on Feb 1, 2023

  1. Merge pull request #2831 from cgarciae:improve-map_variables-v2 

    PiperOrigin-RevId: 506321289
    Flax Authors committed Feb 1, 2023
    f2094ba

  2. Merge pull request #2769 from cgarciae:fix-tabuluate-non-array-values 

    PiperOrigin-RevId: 506321251
    Flax Authors committed Feb 1, 2023
    8c428fa

  3. Update flax/linen/transforms.py

    @chiamp
    chiamp committed Feb 1, 2023
    5f67729

  4. improve the CasualDense example

    @cgarciae
    cgarciae committed Feb 1, 2023
    93e5254

Commits on Jan 31, 2023

  1. import initializers module directly

    @chiamp
    chiamp committed Jan 31, 2023
    d696c70

  2. Merge pull request #2818 from cgarciae:improve-map-variables-docs 

    PiperOrigin-RevId: 506136633
    Flax Authors committed Jan 31, 2023
    e54411d

  3. Merge pull request #2819 from cgarciae:cache-mypy 

    PiperOrigin-RevId: 506121453
    Flax Authors committed Jan 31, 2023
    cf27f3e

  4. Update jax dependency

    @jheek
    jheek committed Jan 31, 2023
    c7fd31f

  5. Add clarification on mixing concrete and abstract vals

    @jheek
    jheek committed Jan 31, 2023
    ed86130
Older