DOCSEARCHbyAlgolia
Submit Docs

sympy2jax Documentation

repository·main·Indexed 18 days ago

https://github.com/patrick-kidger/sympy2jax

sympy2jax turns SymPy expressions into trainable JAX expressions by wrapping them in an Equinox module. It provides the SymbolicModule class to transform PyTrees of SymPy expressions into JAX-compatible modules, allowing for gradient-based optimization of symbolic mathematical models. The library supports custom function mapping via extra_funcs and includes symbolic versions of concatenate and stack operations.

Tokens
2.1K
Snippets
8
Records
9
Agent score
60%

What's inside sympy2jax

  1. Convert SymPy expressions into trainable Equinox modules

    main

    Use sympy2jax.SymbolicModule to transform a PyTree of SymPy expressions into an Equinox module.

    In the resulting module:

    • SymPy symbols become the inputs to the module.
    • SymPy constants (floats, integers, rationals, etc.) become the leaves (parameters) of the module.

    Because the output is an Equinox module, you can optimize these symbolic expressions via gradient descent using standard JAX/Equinox workflows.

    import jax
import sympy
import sympy2jax

x_sym = sympy.symbols("x_sym")
cosx = 1.0 * sympy.cos(x_sym)
sinx = 2.0 * sympy.sin(x_sym)

# Create the module from a list of expressions
mod = sympy2jax.SymbolicModule([cosx, sinx])

# Execute the module by passing symbol names as keyword arguments
x = jax.numpy.zeros(3)
out = mod(x_sym=x)

# Access the trainable parameters (the constants from the SymPy expressions)
params = jax.tree_leaves(mod)  # Returns [1.0, 2.0]

  3. How SymbolicModule works

    main

    SymbolicModule works by transforming a PyTree of SymPy expressions into a PyTree of _AbstractNode objects.

    1. Initialization: During __init__, _sympy_to_node recursively traverses the SymPy expression tree. It maps SymPy atoms (Symbols, Integers, Floats, Constants) and SymPy functions to internal node classes (_Symbol, _Integer, _Float, _Constant, _Func).
    2. Execution: When the module is called (e.g., module(x=1.0)), it uses a memodict to resolve symbols and evaluate the tree. The _Func nodes use a memoization pattern during execution to ensure that intermediate nodes are evaluated efficiently.
    3. Mapping: The core logic relies on a _lookup table that maps SymPy types/functions to JAX/SciPy functions. This mapping includes support for basic arithmetic, trigonometry, exponentials, and linear algebra (like Trace and Determinant).

  4. Use the SymbolicModule API

    main

    The sympy2jax.SymbolicModule class is the primary interface for converting symbolic math to JAX.

    Constructor

    sympy2jax.SymbolicModule(expressions, extra_funcs=None, make_array=True)

    • expressions: A PyTree of SymPy expressions.
    • extra_funcs: (Optional) A dictionary mapping SymPy functions to JAX operations. Use this to extend the built-in translation rules.
    • make_array: (Boolean) If True, integers/floats/rationals are stored as JAX arrays. If False, they are stored as Python scalars.

    Methods

    • __call__(**kwargs): Instances are called using keyword arguments where the keys match the SymPy symbols used in the expressions and the values are the JAX arrays to substitute for those symbols.
    • .sympy(): Translates the module back into a PyTree of SymPy expressions.
    # Example of calling the module
mod = sympy2jax.SymbolicModule([sympy.Symbol('a') * sympy.Symbol('b')])
result = mod(a=1.0, b=2.0)

# Example of converting back to SymPy
sympy_expr = mod.sympy()

  5. Extend SymbolicModule with custom functions

    main

    If SymbolicModule encounters a SymPy function it doesn't recognize, you can extend its capabilities by passing a dictionary to the extra_funcs parameter during initialization. The keys should be the SymPy function objects, and the values should be the corresponding JAX-compatible callables.

    Warning: Providing extra_funcs prevents you from using the .sympy() method to convert the module back to a SymPy expression.

    import sympy as sympy
import jax.numpy as jnp
from sympy2jax import SymbolicModule

x = sympy.Symbol('x')
# Suppose 'special_func' is a SymPy function not in the default lookup
expr = sympy.Function('special_func')(x)

# Map the SymPy function to a JAX implementation
custom_lookup = {sympy.Function('special_func'): jnp.sin}

module = SymbolicModule(expr, extra_funcs=custom_lookup)
result = module(x=0.0)

  6. Use SymbolicModule to convert SymPy expressions to JAX/Equinox

    main

    The SymbolicModule class is the primary interface for converting SymPy expressions (or PyTrees of expressions) into an Equinox module that can be executed using JAX.

    When you initialize a SymbolicModule, it traverses the provided SymPy expressions and builds a tree of internal nodes that map SymPy operations to their JAX equivalents (e.g., sympy.sin becomes jnp.sin).

    Key Features:

    • JAX Compatibility: The resulting module is a valid Equinox module, making it compatible with jax.jit, jax.grad, and other JAX transformations.
    • Custom Functions: You can provide an extra_funcs dictionary to map unsupported SymPy functions to custom JAX functions.
    • Array Conversion: The make_array flag (defaults to True) determines whether constants like integers and floats are converted to JAX arrays.
    • SymPy Reversion: If you did not use extra_funcs, you can call .sympy() to convert the module back into a SymPy expression.
    import sympy as sympy
import equinox as eqx
from sympy2jax import SymbolicModule

# Define a SymPy expression
x = sympy.Symbol('x')
y = sympy.sin(x) + sympy.exp(x)

# Convert to a SymbolicModule
module = SymbolicModule(y)

# Call the module with symbol values
# This returns a JAX array
result = module(x=1.0)

  7. Convert SymbolicModule back to SymPy

    main

    If you initialized a SymbolicModule without providing any extra_funcs, you can convert the module back into a SymPy expression by calling the .sympy() method. This is useful for symbolic manipulation of the resulting JAX-compatible structure.

    Note: This will raise a NotImplementedError if extra_funcs was provided during initialization, as custom functions may not have a defined symbolic representation.

    # Assuming 'module' is a SymbolicModule created without extra_funcs
sympy_expr = module.sympy()

  8. Use SymbolicModule to convert SymPy expressions to JAX

    main

    The SymbolicModule class is the primary interface for converting SymPy expressions into JAX-compatible modules. It is exported from the top-level sympy2jax package. You can use it to wrap symbolic mathematical expressions so they can be used within JAX transformations like jit, grad, or within neural network frameworks like Equinox.

    from sympy2jax import SymbolicModule

# Example usage (conceptual based on API export)
# module = SymbolicModule(sympy_expression, variables)

  9. Use concatenate and stack for symbolic operations

    main

    The sympy2jax package exports concatenate and stack functions. These are intended to provide symbolic versions of JAX's concatenation and stacking operations, allowing you to build complex symbolic expressions that maintain compatibility with JAX's array manipulation patterns.

    from sympy2jax import concatenate, stack