sympy2jax 0.0.6__tar.gz → 0.0.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sympy2jax
-Version: 0.0.6
+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:
@@ -284,6 +284,7 @@ Instances have a `.sympy()` method that translates the module back into a PyTree
 [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.  
 
 **Scientific computing**  
 [Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  

{sympy2jax-0.0.6 → 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:
@@ -61,6 +61,7 @@ Instances have a `.sympy()` method that translates the module back into a PyTree
 [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.  
 
 **Scientific computing**  
 [Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  

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

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

{sympy2jax-0.0.6 → 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
 
 
@@ -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)