sympy2jax 0.0.5__tar.gz → 0.0.7__tar.gz

{sympy2jax-0.0.5 → sympy2jax-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: sympy2jax
-Version: 0.0.5
+Version: 0.0.7
 Summary: Turn SymPy expressions into trainable JAX expressions.
 Project-URL: repository, https://github.com/google/sympy2jax
 Author-email: Patrick Kidger <contact@kidger.site>
@@ -260,7 +260,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:
@@ -274,32 +274,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.  

{sympy2jax-0.0.5 → sympy2jax-0.0.7}/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.  

{sympy2jax-0.0.5 → sympy2jax-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sympy2jax"
-version = "0.0.5"
+version = "0.0.7"
 description = "Turn SymPy expressions into trainable JAX expressions."
 readme = "README.md"
 requires-python ="~=3.9"

{sympy2jax-0.0.5 → sympy2jax-0.0.7}/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)