This library converts SymPy expressions into trainable JAX modules, enabling gradient-based optimization of symbolic mathematical formulas. It is designed for researchers and engineers working with symbolic computation who want to leverage JAX's automatic differentiation and hardware acceleration capabilities for parameter tuning.

How It Works

The core approach involves transforming a PyTree of SymPy expressions into an Equinox module. SymPy symbols are mapped to module inputs, while SymPy numbers (integers, rationals, floats) are treated as trainable parameters (leaves in the PyTree). This allows users to define complex symbolic functions and then optimize their numerical coefficients using standard JAX/Equinox workflows.

Quick Start & Requirements

• Install via pip: pip install sympy2jax
• Requires: Python 3.7+, JAX 0.3.4+, Equinox 0.5.3+, SymPy 1.7.1+.
• Example usage and API documentation are available in the README.

Highlighted Details

• Converts SymPy expressions into Equinox modules for seamless integration with JAX's ecosystem.
• SymPy numbers (integers, rationals, floats) are automatically recognized as trainable parameters.
• Supports custom JAX function mappings for unsupported SymPy functions.
• Can translate optimized modules back into SymPy expressions.

Maintenance & Community

• Developed by Patrick Kidger.
• Links to related libraries in the JAX ecosystem are provided for broader context.

Licensing & Compatibility

• The README does not explicitly state a license. Compatibility for commercial use or closed-source linking is not specified.

Limitations & Caveats

The library is presented as "super easy" with minimal documentation, suggesting it may be best suited for straightforward use cases. The lack of explicit licensing information is a significant caveat for adoption.