Add __eq__ for variables

Author: samanklesaria

flax/nnx/variablelib.py

@@ -777,6 +777,11 @@ def __ror__(self, other) -> A:
       other = other.value
     return self.value.__ror__(other)  # type: ignore
 
+  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

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