Support linen <-> nnx metadata box converging in `nnx.bridge` #4145

IvyZX · 2024-08-27T01:35:13Z

Supports conversion between Linen and NNX metadata boxes.

NNX always has a metadata box (aka nnx.Variable), so make sure ToNNX always make them.
ToLinen will always convert nnx.Variable to a boxed Linen variable. By default, it would be nnx.bridge.NNXMeta, a class that captures everything that an nnx.Variable should.
Already existing boxes of nn.Partitioned should be preserved, and its axes names will be translated to nnx.Variable sharding fields inside ToNNX.

Also refactored variable-related functions that only bridge API uses outside of nnx/variables.py, into a separate nnx/bridge/variables.py.

cgarciae · 2024-08-28T08:52:47Z

flax/nnx/nnx/bridge/variables.py

+
+def register_variable_name_type_pair(name, typ):
+  """Register a pair of variable type name (like Linen collections) and its NNX type."""
+  VariableTypeCache[name] = typ


we should check that the name does exist already

I think we should allow overwrite of names if user wants to. I'll add an overwrite option there.

cgarciae · 2024-08-28T09:08:17Z

flax/nnx/nnx/bridge/variables.py

+  def unbox(self) -> A:
+    return self.value
+
+  def replace_boxed(self, val: B) -> 'NNXMeta[B]':


String types not needed in python >= 3.10

Suggested change

def replace_boxed(self, val: B) -> 'NNXMeta[B]':

def replace_boxed(self, val: B) -> NNXMeta[B]:

I got error trying to do this, saying NNXMeta is not found. I'll leave it as is for now.

cgarciae · 2024-08-28T09:25:17Z

flax/nnx/nnx/bridge/wrappers.py

+
+      nnx_vars = jtu.tree_map_with_path(bv.to_nnx_var, variables,
+                                        is_leaf=lambda x: isinstance(x, meta.AxisMetadata))
+      for col, tree in nnx_vars.items():
        self._setattr(col, tree)


A bit more idiomatic:

Suggested change

self._setattr(col, tree)

setattr(self, col, tree)

cgarciae · 2024-08-28T09:41:54Z

flax/nnx/nnx/bridge/variables.py

+  return NNXMeta(vs.type, vs.value, vs.get_metadata())
+
+
+def to_nnx_var(kp: tp.Sequence[Any], x: meta.AxisMetadata | Any) -> variableslib.Variable:


NIT: We could avoid using key paths here by accepting the collection and iterating over all collections when used.

Suggested change

def to_nnx_var(kp: tp.Sequence[Any], x: meta.AxisMetadata | Any) -> variableslib.Variable:

def to_nnx_var(col: str, x: meta.AxisMetadata | Any) -> variableslib.Variable:

Its a bit more verbose but easier to understand.

Done. We need the key path to get the collection name, so I just moved that part to wrapper.py, to make the function in variables.py cleaner.

cgarciae · 2024-08-28T09:43:17Z

flax/nnx/nnx/bridge/wrappers.py

      for collection, value in updates.items():
-        self._setattr(collection, jax.tree.map(variableslib.variable_type(collection), value))
+        self._setattr(collection, value)


Suggested change

self._setattr(collection, value)

setattr(self, collection, value)

cgarciae · 2024-08-28T13:45:36Z

flax/nnx/nnx/bridge/wrappers.py

+
+      nnx_vars = jtu.tree_map_with_path(bv.to_nnx_var, variables,
+                                        is_leaf=lambda x: isinstance(x, meta.AxisMetadata))
+      for col, tree in nnx_vars.items():
        self._setattr(col, tree)
        self.linen_collections.add(col)


For linen_collections we should either implement this with a list or create a PytreeSet. The issue is that sets are not pytrees so they are treated as static but they are not hashable.

list is mutable so not hashable either. Curious why that matters though... but I'll make it a tuple.

cgarciae · 2024-08-28T14:26:28Z

flax/nnx/nnx/bridge/wrappers.py

@@ -225,14 +230,16 @@ def update_variables(self, module):
    types = set(jax.tree.leaves(
      jax.tree.map(lambda x: x.type, state,
                    is_leaf=lambda x: isinstance(x, nnx.VariableState))))
-    types = variableslib.sort_variable_types(types)
+    types = bv.sort_variable_types(types)
    _, *state_by_types = nnx.split(module, *types)


Use State.split:

Suggested change

_, *state_by_types = nnx.split(module, *types)

*state_by_types = state.split(*types)

Thanks but nnx.split is more convenient because it always return a tuple...

cgarciae · 2024-08-28T14:29:59Z

flax/nnx/nnx/bridge/wrappers.py

@@ -225,14 +230,16 @@ def update_variables(self, module):
    types = set(jax.tree.leaves(
      jax.tree.map(lambda x: x.type, state,
                    is_leaf=lambda x: isinstance(x, nnx.VariableState))))
-    types = variableslib.sort_variable_types(types)
+    types = bv.sort_variable_types(types)


sort_variable_types expects type list but we can change the type to Iterable

Thanks for the catch! Done.

IvyZX requested review from levskaya and cgarciae August 27, 2024 01:35

IvyZX force-pushed the bdg branch 2 times, most recently from aa7f87a to de17c01 Compare August 27, 2024 21:58

cgarciae reviewed Aug 28, 2024

View reviewed changes

cgarciae approved these changes Aug 28, 2024

View reviewed changes

IvyZX force-pushed the bdg branch from de17c01 to 450e80b Compare August 28, 2024 20:34

Support linen <-> nnx metadata box converging in bridge

4bb2152

IvyZX force-pushed the bdg branch from 450e80b to 4bb2152 Compare August 28, 2024 20:49

IvyZX added the pull ready label Aug 28, 2024

copybara-service bot merged commit 839db8c into google:main Aug 28, 2024
16 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support linen <-> nnx metadata box converging in `nnx.bridge` #4145

Support linen <-> nnx metadata box converging in `nnx.bridge` #4145

IvyZX commented Aug 27, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

cgarciae Aug 28, 2024

IvyZX Aug 28, 2024

	def replace_boxed(self, val: B) -> 'NNXMeta[B]':
	def replace_boxed(self, val: B) -> NNXMeta[B]:

		return NNXMeta(vs.type, vs.value, vs.get_metadata())


		def to_nnx_var(kp: tp.Sequence[Any], x: meta.AxisMetadata \| Any) -> variableslib.Variable:

	def to_nnx_var(kp: tp.Sequence[Any], x: meta.AxisMetadata \| Any) -> variableslib.Variable:
	def to_nnx_var(col: str, x: meta.AxisMetadata \| Any) -> variableslib.Variable:

	self._setattr(collection, value)
	setattr(self, collection, value)

	_, state_by_types = nnx.split(module, types)
	state_by_types = state.split(types)

Support linen <-> nnx metadata box converging in nnx.bridge #4145

Support linen <-> nnx metadata box converging in nnx.bridge #4145

Conversation

IvyZX commented Aug 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Support linen <-> nnx metadata box converging in `nnx.bridge` #4145

Support linen <-> nnx metadata box converging in `nnx.bridge` #4145