sympy2jax 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

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)

{sympy2jax-0.0.5.dist-info → sympy2jax-0.0.7.dist-info}/METADATA

@@ -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.7.dist-info/RECORD

@@ -0,0 +1,6 @@
+sympy2jax/__init__.py,sha256=KqjLNlIiDATWsPhjPLpG-Hud4HMMuq_DP6sISDWVYQg,720
+sympy2jax/sympy_module.py,sha256=R1RjJ57jvLcSqLJZNu20BE3XdYHMi2UuNGmZslifYyA,9411
+sympy2jax-0.0.7.dist-info/METADATA,sha256=lkRfJ18NFu93bdodMkm7DrqLYYLayJh-qwce0rR5C8o,16659
+sympy2jax-0.0.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sympy2jax-0.0.7.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+sympy2jax-0.0.7.dist-info/RECORD,,

{sympy2jax-0.0.5.dist-info → sympy2jax-0.0.7.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.18.0
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any

sympy2jax-0.0.5.dist-info/RECORD

@@ -1,6 +0,0 @@
-sympy2jax/__init__.py,sha256=KqjLNlIiDATWsPhjPLpG-Hud4HMMuq_DP6sISDWVYQg,720
-sympy2jax/sympy_module.py,sha256=t1Z8yTRCDgDY_CcG8TF19aAJivKSdz__Ej_OdQ0272g,9378
-sympy2jax-0.0.5.dist-info/METADATA,sha256=iJitCmILWFrGKzgHMNxSIUKeX20FrD4alddTdNNHrKc,16459
-sympy2jax-0.0.5.dist-info/WHEEL,sha256=9QBuHhg6FNW7lppboF2vKVbCGTVzsFykgRQjjlajrhA,87
-sympy2jax-0.0.5.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-sympy2jax-0.0.5.dist-info/RECORD,,