[tree-math] add replace() method

This convenience method is copied from flax.struct.

PiperOrigin-RevId: 621690005

setup.py

@@ -29,7 +29,7 @@
 setuptools.setup(
     name='tree-math',
     description='Mathematical operations for JAX pytrees',
-    version='0.2.0 ',
+    version='0.2.1',
     license='Apache 2.0',
     author='Google LLC',
     author_email='[email protected]',

tree_math/__init__.py

@@ -24,4 +24,4 @@
 from tree_math._src.vector import Vector, VectorMixin
 import tree_math.numpy
 
-__version__ = '0.2.0'
+__version__ = '0.2.1'

tree_math/_src/structs.py

@@ -72,6 +72,7 @@ def tree_unflatten(cls, _, children):
                        {'fields': fields,
                         'asdict': asdict,
                         'astuple': astuple,
+                        'replace': dataclasses.replace,
                         'tree_flatten': tree_flatten,
                         'tree_unflatten': tree_unflatten,
                         '__module__': cls.__module__})

tree_math/_src/structs_test.py

@@ -109,6 +109,12 @@ def testPickle(self):
     restored = pickle.loads(pickle.dumps(struct))
     self.assertTreeEqual(struct, restored, check_dtypes=True)
 
+  def testReplace(self):
+    struct = TestStruct(1, 2)
+    replaced = struct.replace(b=3)
+    expected = TestStruct(1, 3)
+    self.assertTreeEqual(replaced, expected, check_dtypes=True)
+
 
 if __name__ == '__main__':
   absltest.main()