Fixed .type deprecated property usage

Author: vfdev-5

flax/nnx/bridge/variables.py

@@ -97,7 +97,7 @@ def to_linen_var(vs: variablelib.Variable) -> meta.AxisMetadata:
     return linen_type(vs.value, **metadata)
   if is_vanilla_variable(vs):
     return vs.value
-  return NNXMeta(vs.type, vs.value, metadata)
+  return NNXMeta(type(vs), vs.value, metadata)
 
 
 def get_col_name(keypath: tp.Sequence[Any]) -> str:

flax/nnx/bridge/wrappers.py

@@ -349,7 +349,7 @@ def _update_variables(self, module):
 
     # group state by collection
     for path, leaf in nnx.to_flat_state(state):
-      type_ = leaf.type if isinstance(leaf, nnx.Variable) else type(leaf)
+      type_ = type(leaf)
       collection = variablelib.variable_name_from_type(
           type_, allow_register=True
       )

tests/nnx/transforms_test.py

@@ -628,9 +628,9 @@ def f(m: Dict):
     assert m.a[0] is m.b
     assert isinstance(grads, nnx.State)
     assert grads['a']['0'].value == 2.0
-    assert issubclass(grads['a']['0'].type, nnx.Variable)
+    assert issubclass(type(grads['a']['0']), nnx.Variable)
     assert grads['a']['1'].value == 1.0
-    assert issubclass(grads['a']['1'].type, nnx.Variable)
+    assert issubclass(type(grads['a']['1']), nnx.Variable)
     assert len(nnx.to_flat_state(grads)) == 2
 
     nnx.update(m, grads)
@@ -659,7 +659,7 @@ def f(m: Dict):
 
     assert isinstance(grads, nnx.State)
     assert grads['a']['0'].value == 1.0
-    assert issubclass(grads['a']['0'].type, nnx.Param)
+    assert issubclass(type(grads['a']['0']), nnx.Param)
     assert len(grads) == 2
 
     nnx.update(m, grads)
@@ -687,7 +687,7 @@ def f(m: Dict):
 
     assert isinstance(grads, nnx.State)
     assert grads['a']['1'].value == 1.0
-    assert issubclass(grads['a']['1'].type, nnx.BatchStat)
+    assert issubclass(type(grads['a']['1']), nnx.BatchStat)
     assert len(grads) == 1
 
     nnx.update(m, grads)
@@ -843,9 +843,9 @@ def f(m: dict):
 
     assert m['a'][0] is m['b']
     assert isinstance(grads, dict)
-    assert issubclass(grads['a'][0].type, nnx.Variable)
+    assert issubclass(type(grads['a'][0]), nnx.Variable)
     assert grads['a'][1].value == 1.0
-    assert issubclass(grads['a'][1].type, nnx.Variable)
+    assert issubclass(type(grads['a'][1]), nnx.Variable)
     assert len(jax.tree.leaves(grads)) == 2
 
     jax.tree.map(