code: def test_magic_methods(self):
      def same_value(v1, v2):
          if isinstance(v1, tuple):
              return all(
                  [
                      jnp.all(jax.tree_map(lambda x, y: x == y, e1, e2))
                      for e1, e2 in zip(v1, v2)
                  ]
              )
          return jnp.all(jax.tree_map(lambda x, y: x == y, v1, v2))
  
      list_obj = [(1, 2), (3, 4)]
      for l, tuple_wrap in (
          (list_obj, lambda x: x),
          (tuple(list_obj), tuple),
      ):
          c = cursor(l)
          # test __len__
          self.assertTrue(same_value(len(c), len(l)))
          # test __iter__
          for i, child_c in enumerate(c):
              child_c[1] += i + 1
          self.assertEqual(c.build(), tuple_wrap([(1, 3), (3, 6)]))
          # test __reversed__
          for i, child_c in enumerate(reversed(c)):
              child_c[1] += i + 1
          self.assertEqual(c.build(), tuple_wrap([(1, 5), (3, 7)]))
      # test __iter__ error
      with self.assertRaisesRegex(
          NotImplementedError,
          "__iter__ method only implemented for tuples and lists, not type <class"
          " 'dict'>",
      ):
          c = cursor({"a": 1, "b": 2})
          for key in c:
              c[key] *= -1
      # test __iter__ error
      with self.assertRaisesRegex(
          NotImplementedError,
          "__reversed__ method only implemented for tuples and lists, not type"
          " <class 'dict'>",
      ):
          c = cursor({"a": 1, "b": 2})
          for key in reversed(c):
              c[key] *= -1
  
      for obj_value in (2, jnp.array([[1, -2], [3, 4]])):
          for c in (
              cursor(obj_value),
              cursor([obj_value])[0],
              cursor((obj_value,))[0],
              cursor({0: obj_value})[0],
              cursor(flax.core.freeze({0: obj_value}))[0],
              cursor(GenericTuple(x=obj_value)).x,
              cursor(GenericDataClass(x=obj_value)).x,
          ):
              # test __neg__
              self.assertTrue(same_value(-c, -obj_value))
              # test __pos__
              self.assertTrue(same_value(+c, +obj_value))
              # test __abs__
              self.assertTrue(same_value(abs(-c), abs(-obj_value)))
              # test __invert__
              self.assertTrue(same_value(~c, ~obj_value))
              # test __round__
              self.assertTrue(
                  same_value(round(c + 0.123), round(obj_value + 0.123))
              )
              self.assertTrue(
                  same_value(
                      round(c + 0.123, 2), round(obj_value + 0.123, 2)
                  )
              )
  
              for other_value in (3, jnp.array([[5, 6], [7, 8]])):
                  # test __add__
                  self.assertTrue(
                      same_value(
                          c + other_value, obj_value + other_value
                      )
                  )
                  # test __radd__
                  self.assertTrue(
                      same_value(
                          other_value + c, other_value + obj_value
                      )
                  )
                  # test __sub__
                  self.assertTrue(
                      same_value(
                          c - other_value, obj_value - other_value
                      )
                  )
                  # test __rsub__
                  self.assertTrue(
                      same_value(
                          other_value - c, other_value - obj_value
                      )
                  )
                  # test __mul__
                  self.assertTrue(
                      same_value(
                          c * other_value, obj_value * other_value
                      )
                  )
                  # test __rmul__
                  self.assertTrue(
                      same_value(
                          other_value * c, other_value * obj_value
                      )
                  )
                  # test __truediv__
                  self.assertTrue(
                      same_value(
                          c / other_value, obj_value / other_value
                      )
                  )
                  # test __rtruediv__
                  self.assertTrue(
                      same_value(
                          other_value / c, other_value / obj_value
                      )
                  )
                  # test __floordiv__
                  self.assertTrue(
                      same_value(
                          c // other_value, obj_value // other_value
                      )
                  )
                  # test __rfloordiv__
                  self.assertTrue(
                      same_value(
                          other_value // c, other_value // obj_value
                      )
                  )
                  # test __mod__
                  self.assertTrue(
                      same_value(
                          c % other_value, obj_value % other_value
                      )
                  )
                  # test __rmod__
                  self.assertTrue(
                      same_value(
                          other_value % c, other_value % obj_value
                      )
                  )
                  # test __divmod__
                  self.assertTrue(
                      same_value(
                          divmod(c, other_value),
                          divmod(obj_value, other_value),
                      )
                  )
                  # test __rdivmod__
                  self.assertTrue(
                      same_value(
                          divmod(other_value, c),
                          divmod(other_value, obj_value),
                      )
                  )
                  # test __pow__
                  self.assertTrue(
                      same_value(
                          pow(c, other_value),
                          pow(obj_value, other_value),
                      )
                  )
                  # test __rpow__
                  self.assertTrue(
                      same_value(
                          pow(other_value, c),
                          pow(other_value, obj_value),
                      )
                  )
                  # test __lshift__
                  self.assertTrue(
                      same_value(
                          c << other_value, obj_value << other_value
                      )
                  )
                  # test __rlshift__
                  self.assertTrue(
                      same_value(
                          other_value << c, other_value << obj_value
                      )
                  )
                  # test __rshift__
                  self.assertTrue(
                      same_value(
                          c >> other_value, obj_value >> other_value
                      )
                  )
                  # test __rrshift__
                  self.assertTrue(
                      same_value(
                          other_value >> c, other_value >> obj_value
                      )
                  )
                  # test __and__
                  self.assertTrue(
                      same_value(
                          c & other_value, obj_value & other_value
                      )
                  )
                  # test __rand__
                  self.assertTrue(
                      same_value(
                          other_value & c, other_value & obj_value
                      )
                  )
                  # test __xor__
                  self.assertTrue(
                      same_value(
                          c ^ other_value, obj_value ^ other_value
                      )
                  )
                  # test __rxor__
                  self.assertTrue(
                      same_value(
                          other_value ^ c, other_value ^ obj_value
                      )
                  )
                  # test __or__
                  self.assertTrue(
                      same_value(
                          c | other_value, obj_value | other_value
                      )
                  )
                  # test __ror__
                  self.assertTrue(
                      same_value(
                          other_value | c, other_value | obj_value
                      )
                  )
  
                  if isinstance(obj_value, jax.Array) and isinstance(
                      other_value, jax.Array
                  ):
                      # test __matmul__
                      self.assertTrue(
                          same_value(
                              c @ other_value, obj_value @ other_value
                          )
                      )
                      # test __rmatmul__
                      self.assertTrue(
                          same_value(
                              other_value @ c, other_value @ obj_value
                          )
                      )
  
                  # test __lt__
                  self.assertTrue(
                      same_value(
                          c < other_value, obj_value < other_value
                      )
                  )
                  self.assertTrue(
                      same_value(
                          other_value < c, other_value < obj_value
                      )
                  )
                  # test __le__
                  self.assertTrue(
                      same_value(
                          c <= other_value, obj_value <= other_value
                      )
                  )
                  self.assertTrue(
                      same_value(
                          other_value <= c, other_value <= obj_value
                      )
                  )
                  # test __eq__
                  self.assertTrue(
                      same_value(
                          c == other_value, obj_value == other_value
                      )
                  )
                  self.assertTrue(
                      same_value(
                          other_value == c, other_value == obj_value
                      )
                  )
                  # test __ne__
                  self.assertTrue(
                      same_value(
                          c != other_value, obj_value != other_value
                      )
                  )
                  self.assertTrue(
                      same_value(
                          other_value != c, other_value != obj_value
                      )
                  )
                  # test __gt__
                  self.assertTrue(
                      same_value(
                          c > other_value, obj_value > other_value
                      )
                  )
                  self.assertTrue(
                      same_value(
                          other_value > c, other_value > obj_value
                      )
                  )
                  # test __ge__
                  self.assertTrue(
                      same_value(
                          c >= other_value, obj_value >= other_value
                      )
                  )
                  self.assertTrue(
                      same_value(
                          other_value >= c, other_value >= obj_value
                      )
                  )
  
imports: import dataclasses
  from typing import Any, NamedTuple
  import jax
  import jax.numpy
  import optax
  from absl.testing import absltest
  import flax
  import flax.linen
  from flax.core import freeze
  from flax.cursor import AccessType, _traverse_tree, cursor
  from flax.errors import CursorFindError, TraverseTreeError
  from flax.training import train_state
label_desc: Use image node_id 2 for calling the CursorTest obj's underlying member method code with example usage: obj.test_magic_methods() and returns: jnp, all