Merge pull request #33 from leakec/jax_7_1

Updating for JAX 0.7.1

Author: leakec

README.md

@@ -66,7 +66,7 @@ The authors of this repsitory and the associated theory have gone to lengths to
     author = {Carl Leake and Hunter Johnston},
     title = {{TFC: A Functional Interpolation Framework}},
     url = {https://github.com/leakec/tfc},
-    version = {1.3.0},
+    version = {1.3.1},
     year = {2025},
 }
 @article{TFC, 

pyproject.toml

@@ -8,8 +8,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tfc"
-version = "1.3.0"
-requires-python = ">=3.11,<3.14"
+version = "1.3.1"
+requires-python = ">=3.11"
 readme = "README.md"
 dynamic = ["dependencies", "classifiers", "authors", "license", "description"]
 

src/tfc/mtfc.py

@@ -623,10 +623,7 @@ def H_jvp(arg_vals, arg_tans, d: tuple[int, ...] = d0, full: bool = False):
             out_tans = np.zeros((dim0, dim1))
             for k in range(n):
                 if not (type(arg_tans[k]) is ad.Zero):
-                    if type(arg_tans[k]) is batching.BatchTracer:
-                        flag = onp.any(arg_tans[k].val != 0)
-                    else:
-                        flag = onp.any(arg_tans[k] != 0)
+                    flag = onp.any(arg_tans[k] != 0)
                     if flag:
                         dark = tuple(d[j] + 1 if k == j else d[j] for j in range(len(d)))
                         if flat:

src/tfc/utfc.py

@@ -362,10 +362,7 @@ def H_jvp(arg_vals, arg_tans, d: uint = 0, full: bool = False):
             x = arg_vals[0]
             dx = arg_tans[0]
             if not (dx is ad.Zero):
-                if type(dx) is batching.BatchTracer:
-                    flag = onp.any(dx.val != 0)
-                else:
-                    flag = onp.any(dx != 0)
+                flag = onp.any(dx != 0)
                 if flag:
                     if len(dx.shape) == 1:
                         out_tans = Hjax(x, d=d + 1, full=full) * onp.expand_dims(dx, 1)

src/tfc/version.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.0"
+__version__ = "1.3.1"
 
 
 def _version_as_tuple(version_str: str) -> tuple[int, ...]: