Repository: patrick-kidger/sympy2jax

Issue:
DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
Using:
```toml
[[package]]
name = "jax"
version = "0.4.26"

[[package]]
name = "equinox"
version = "0.11.3"
```

```
../../.cache/pypoetry/virtualenvs/matrixelement-WHObBYqn-py3.11/lib/python3.11/site-packages/sympy2jax/sympy_module.py:310: in __init__
    self.nodes = jax.tree_map(_convert, expressions)
```

Discussion:

Solution patch:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index a4e560c..fc37397 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Release
-        uses: patrick-kidger/action_update_python_project@v2
+        uses: patrick-kidger/action_update_python_project@v6
         with:
             python-version: "3.11"
             test-script: |
@@ -20,7 +20,3 @@ jobs:
             pypi-token: ${{ secrets.pypi_token }}
             github-user: patrick-kidger
             github-token: ${{ github.token }}
-            email-user: ${{ secrets.email_user }}
-            email-token: ${{ secrets.email_token }}
-            email-server: ${{ secrets.email_server }}
-            email-target: ${{ secrets.email_target }}
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
index 59d9ebd..a2f50e5 100644
--- a/.github/workflows/run_tests.yml
+++ b/.github/workflows/run_tests.yml
@@ -7,7 +7,7 @@ jobs:
   run-tests:
     strategy:
       matrix:
-        python-version: [ 3.9 ]
+        python-version: [ "3.10", "3.12" ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b6033ae..24a87e0 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,23 +1,14 @@
 repos:
-  - repo: https://github.com/ambv/black
-    rev: 22.3.0
-    hooks:
-    - id: black
-  - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.0.255'
-    hooks:
-      - id: ruff
-        args: ["--fix"]
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.2
+    hooks:
+      - id: ruff-format  # formatter
+        types_or: [ python, pyi, jupyter ]
+      - id: ruff  # linter
+        types_or: [ python, pyi, jupyter ]
+        args: [ --fix ]
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.315
+    rev: v1.1.350
     hooks:
     - id: pyright
-      additional_dependencies: ["equinox", "jax", "sympy"]
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.6.3
-    hooks:
-    - id: nbqa-black
-      additional_dependencies: [ipython==8.12, black]
-    - id: nbqa-ruff 
-      args: ["--ignore=I001"]
-      additional_dependencies: [ipython==8.12, ruff]
+      additional_dependencies: [equinox, jax, sympy]
diff --git a/README.md b/README.md
index 26b99ef..a4231e3 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,7 @@ params = jax.tree_leaves(mod)  # 1.0 and 2.0 are parameters.
 ## Documentation
 
 ```python
-sympytorch.SymbolicModule(expressions, extra_funcs=None, make_array=True)
+sympy2jax.SymbolicModule(expressions, extra_funcs=None, make_array=True)
 ```
 
 Where:
@@ -51,32 +51,24 @@ Instances have a `.sympy()` method that translates the module back into a PyTree
 
 (That's literally the entire documentation, it's super easy.)
 
-## Finally
+## See also: other libraries in the JAX ecosystem
 
-### See also: other libraries in the JAX ecosystem
+**Always useful**  
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks and everything not already in core JAX!  
+[jaxtyping](https://github.com/patrick-kidger/jaxtyping): type annotations for shape/dtype of arrays.  
 
-[jaxtyping](https://github.com/google/jaxtyping): type annotations for shape/dtype of arrays.
+**Deep learning**  
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.  
+[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).  
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).  
+[paramax](https://github.com/danielward27/paramax): parameterizations and constraints for PyTrees.  
 
-[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
+**Scientific computing**  
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  
+[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.  
+[Lineax](https://github.com/patrick-kidger/lineax): linear solvers.  
+[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.  
+[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  
 
-[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
-
-[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
-
-[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.
-
-[Lineax](https://github.com/google/lineax): linear solvers.
-
-[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.
-
-[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).
-
-[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
-
-[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).
-
-[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)
-
-### Disclaimer
-
-This is not an official Google product.
+**Awesome JAX**  
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  
diff --git a/pyproject.toml b/pyproject.toml
index f6db440..c8624a9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "sympy2jax"
-version = "0.0.5"
+version = "0.0.6"
 description = "Turn SymPy expressions into trainable JAX expressions."
 readme = "README.md"
 requires-python ="~=3.9"
diff --git a/sympy2jax/sympy_module.py b/sympy2jax/sympy_module.py
index 7d5d34e..47adbc9 100644
--- a/sympy2jax/sympy_module.py
+++ b/sympy2jax/sympy_module.py
@@ -22,6 +22,7 @@ import equinox as eqx
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
+import jax.tree_util as jtu
 import sympy
 
 
@@ -129,7 +130,7 @@ class _Symbol(_AbstractNode):
     _name: str
 
     def __init__(self, expr: sympy.Expr):
-        self._name = expr.name  # pyright: ignore
+        self._name = str(expr.name)  # pyright: ignore
 
     def __call__(self, memodict: dict):
         try:
@@ -307,7 +308,7 @@ class SymbolicModule(eqx.Module):
             func_lookup=lookup,
             make_array=make_array,
         )
-        self.nodes = jax.tree_map(_convert, expressions)
+        self.nodes = jtu.tree_map(_convert, expressions)
 
     def sympy(self) -> sympy.Expr:
         if self.has_extra_funcs:
@@ -316,10 +317,10 @@ class SymbolicModule(eqx.Module):
                 "is passed."
             )
         memodict = dict()
-        return jax.tree_map(
+        return jtu.tree_map(
             lambda n: n.sympy(memodict, _reverse_lookup), self.nodes, is_leaf=_is_node
         )
 
     def __call__(self, **symbols):
         memodict = symbols
-        return jax.tree_map(lambda n: n(memodict), self.nodes, is_leaf=_is_node)
+        return jtu.tree_map(lambda n: n(memodict), self.nodes, is_leaf=_is_node)


Test patch:
diff --git a/tests/test_symbolic_module.py b/tests/test_symbolic_module.py
index 6aff574..755a331 100644
--- a/tests/test_symbolic_module.py
+++ b/tests/test_symbolic_module.py
@@ -57,8 +57,8 @@ def assert_sympy_allclose(x, y):
 
 def test_example():
     x_sym = sympy.symbols("x_sym")
-    cosx = 1.0 * sympy.cos(x_sym)
-    sinx = 2.0 * sympy.sin(x_sym)
+    cosx = 1.0 * sympy.cos(x_sym)  # pyright: ignore[reportOperatorIssue]
+    sinx = 2.0 * sympy.sin(x_sym)  # pyright: ignore[reportOperatorIssue]
     mod = sympy2jax.SymbolicModule([cosx, sinx])
 
     x = jax.numpy.zeros(3)