Fixed deprecated .value usage failing CI tests

Author: vfdev-5

flax/nnx/variablelib.py

@@ -1712,25 +1712,25 @@ def __contains__(self, item) -> bool:
 
   def __eq__(self, other) -> bool:
     if isinstance(other, Variable):
-      other = other.value
-    return self.value.__eq__(other)  # type: ignore
+      other = other.get_value()
+    return self.get_value().__eq__(other)  # type: ignore
 
   def __iadd__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value += x` instead.'
+      'Use `variable[...] += x` instead.'
     )
 
   def __isub__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value -= x` instead.'
+      'Use `variable[...] -= x` instead.'
     )
 
   def __imul__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value *= x` instead.'
+      'Use `variable[...] *= x` instead.'
     )
 
   def __imatmul__(self: V, other) -> V:

tests/nnx/variable_test.py

@@ -18,12 +18,12 @@
 import jax.numpy as jnp
 import pytest
 
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 from flax import nnx
 
 A = tp.TypeVar('A')
 
-class TestVariable(absltest.TestCase):
+class TestVariable(parameterized.TestCase):
   def test_pytree(self):
     r1 = nnx.Param(1)
     self.assertEqual(r1.get_value(), 1)
@@ -95,6 +95,18 @@ def test_binary_ops(self):
 
     self.assertEqual(v1[...], 5)
 
+  @parameterized.product(
+    v1=[jnp.array([1, 2]), jnp.array(2), 3],
+    v2=[jnp.array([1, 2]), jnp.array(2), 3],
+  )
+  def test_eq_op(self, v1, v2):
+    p1 = nnx.Param(v1)
+    p2 = nnx.Param(v2)
+    if isinstance(v1, jax.Array) or isinstance(v2, jax.Array):
+      self.assertEqual((p1 == p2).all(), (v1 == v2).all())
+    else:
+      self.assertEqual(p1 == p2, v1 == v2)
+
   def test_mutable_array_context(self):
     initial_mode = nnx.using_hijax()
     with nnx.use_hijax(False):