Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[nnx] remove pytreelib #3816

Merged
merged 1 commit into from
Apr 10, 2024
Merged

[nnx] remove pytreelib #3816

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions flax/experimental/nnx/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,6 @@
from .nnx.nn.normalization import LayerNorm as LayerNorm
from .nnx.nn.normalization import RMSNorm as RMSNorm
from .nnx.nn.stochastic import Dropout as Dropout
from .nnx.pytreelib import Pytree as Pytree
from .nnx.pytreelib import TreeNode as TreeNode
from .nnx.rnglib import Rngs as Rngs
from .nnx.rnglib import RngStream as RngStream
from .nnx.rnglib import RngState as RngState
Expand Down
269 changes: 0 additions & 269 deletions flax/experimental/nnx/nnx/pytreelib.py
View file
Open in desktop

This file was deleted.

12 changes: 5 additions & 7 deletions flax/experimental/nnx/nnx/spmd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from jax.sharding import Mesh, PartitionSpec

from flax.experimental.nnx.nnx import variables
from flax.experimental.nnx.nnx.pytreelib import TreeNode
from flax.experimental.nnx.nnx.state import State
from flax.typing import (
Array,
Expand Down Expand Up @@ -107,17 +106,16 @@ def f(x):

return _maybe_replicate(x)

return jax.tree_util.tree_map(
f,
tree,
is_leaf=lambda x: isinstance(x, variables.Variable)
and not isinstance(x, TreeNode),
return jax.tree_map(
f, tree, is_leaf=lambda x: isinstance(x, variables.Variable)
)


def get_named_sharding(tree: A, mesh: jax.sharding.Mesh) -> A:
spec = get_partition_spec(tree)
sharding = jax.tree_util.tree_map(lambda p: jax.sharding.NamedSharding(mesh, p), spec)
sharding = jax.tree_util.tree_map(
lambda p: jax.sharding.NamedSharding(mesh, p), spec
)
return sharding


Expand Down
6 changes: 3 additions & 3 deletions flax/experimental/nnx/tests/test_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -527,15 +527,15 @@ def test_basic(self):
@dataclasses.dataclass
class Foo(nnx.Module):
a: int
b: nnx.TreeNode[int]
b: nnx.Variable[int]
c: nnx.Param[int]
d: nnx.Variable[int]
e: nnx.Variable[int]
f: int

m = Foo(
a=1, # static
b=nnx.TreeNode(2), # node
b=nnx.Variable(2), # node
c=nnx.Param(3), # param
d=nnx.Variable(4), # var
e=nnx.BatchStat(5), # var
Expand All @@ -545,7 +545,7 @@ class Foo(nnx.Module):
graphdef, state = m.split()

assert len(state) == 4
assert state.b == nnx.TreeNode(2)
assert state.b == nnx.Variable(2)
assert state.c == nnx.Param(3)
assert state.d == nnx.Variable(4)
assert state.e == nnx.BatchStat(5)
Expand Down
Loading
Loading