Merge pull request #5084 from samanklesaria:variable_eq

PiperOrigin-RevId: 833483732

Author: Flax Authors

flax/nnx/variablelib.py

@@ -1710,6 +1710,11 @@ def __contains__(self, item) -> bool:
   __rxor__ = _variable_operator('__rxor__')
   __ror__ = _variable_operator('__ror__')
 
+  def __eq__(self, other) -> bool:
+    if isinstance(other, Variable):
+      other = other.value
+    return self.value.__eq__(other)  # type: ignore
+
   def __iadd__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'

tests/nnx/variable_test.py

@@ -89,6 +89,7 @@ def test_binary_ops(self):
     result = v1 + v2
 
     self.assertEqual(result, 5)
+    self.assertFalse(v1 == v2)
 
     v1[...] += v2