Add flax_field and tree_map_with_path

tjax/__init__.py

@@ -16,14 +16,16 @@
                                    replace_cotangent, scale_cotangent)
 from ._src.display import display_generic, print_generic, tapped_print_generic
 from ._src.dtypes import default_atol, default_rtol, default_tols
+from ._src.flax_tools import flax_field
 from ._src.leaky_integral import (diffused_leaky_integrate, leaky_covariance, leaky_data_weight,
                                   leaky_integrate, leaky_integrate_time_series)
+from ._src.math_tools import (abs_square, divide_nonnegative, divide_where, inverse_softplus,
+                              is_scalar, zero_tangent_like)
 from ._src.partial import Partial
 from ._src.shims import custom_jvp, custom_jvp_method, custom_vjp, custom_vjp_method, jit
 from ._src.testing import (assert_tree_allclose, get_relative_test_string, get_test_string,
                            tree_allclose)
-from ._src.tools import (abs_square, divide_nonnegative, divide_where, inverse_softplus, is_scalar,
-                         zero_tangent_like)
+from ._src.tree_tools import tree_map_with_path
 
 __all__ = ['BooleanArray', 'BooleanNumeric', 'Complex', 'ComplexArray', 'ComplexNumeric', 'Array',
            'KeyArray', 'Integral', 'IntegralArray', 'IntegralNumeric', 'NumpyArray',
@@ -40,7 +42,7 @@
            'cotangent_combinator', 'print_cotangent', 'print_generic', 'replace_cotangent',
            'JaxArray', 'JaxBooleanArray', 'JaxIntegralArray', 'tree_allclose', 'zero_tangent_like',
            'JaxAbstractClass', 'JaxComplexArray', 'JaxRealArray', 'abstract_jit',
-           'abstract_custom_jvp']
+           'abstract_custom_jvp', 'tree_map_with_path', 'flax_field']
 #
 # __pdoc__ = {}
 # __pdoc__['PyTreeLike'] = False

tjax/_src/__init__.py

@@ -6,7 +6,7 @@
 from .dtypes import *
 from .graph import *
 from .leaky_integral import *
+from .math_tools import *
 from .partial import *
 from .shims import *
 from .testing import *
-from .tools import *

tjax/_src/fixed_point/comparing.py

@@ -9,7 +9,7 @@
 
 from ..annotations import JaxBooleanArray, JaxRealArray
 from ..dataclasses import dataclass
-from ..tools import divide_nonnegative
+from ..math_tools import divide_nonnegative
 from .augmented import AugmentedState, State
 from .iterated_function import Comparand, IteratedFunction, Parameters, Trajectory
 

tjax/_src/fixed_point/stochastic.py

@@ -10,7 +10,7 @@
 from ..annotations import ComplexNumeric, JaxBooleanArray, JaxComplexArray, JaxRealArray
 from ..dataclasses import dataclass
 from ..leaky_integral import leaky_data_weight, leaky_integrate
-from ..tools import abs_square, divide_nonnegative
+from ..math_tools import abs_square, divide_nonnegative
 from .augmented import AugmentedState, State
 from .combinator import Differentiand, IteratedFunctionWithCombinator
 from .iterated_function import Comparand, IteratedFunction, Parameters, Trajectory

tjax/_src/flax_tools.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from dataclasses import field
+from typing import Any
+
+__all__ = ['flax_field']
+
+
+def flax_field() -> Any:
+    return field(init=False, default=None, kw_only=True)

tjax/_src/gradient/transform.py

@@ -9,7 +9,7 @@
 
 from ..annotations import PyTree
 from ..dataclasses import dataclass
-from ..tools import abs_square
+from ..math_tools import abs_square
 
 __all__ = ['GradientState', 'GenericGradientState', 'GradientTransformation',
            'SecondOrderGradientTransformation', 'ThirdOrderGradientTransformation']

tjax/_src/tree_tools.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+__all__ = ['tree_map_with_path']
+
+
+T = TypeVar('T')
+def tree_map_with_path(structure: Any,
+                       transform: Callable[[T, tuple[str, ...]], T]
+                       ) -> Any:
+    def f(structure: Any,
+          transform: Callable[[T, tuple[str, ...]], T],
+          path: tuple[str, ...]
+          ) -> Any:
+        if isinstance(structure, dict):
+            out_structure = {}
+            for key, value in structure.items():
+                out_structure[key] = f(value, transform, (*path, key))
+            return out_structure
+        return transform(structure, path)
+    return f(structure, transform, ())