Turning SymPy expressions into JAX functions

Last update: Dec 11, 2022

Related tags

Deep Learning sympy2jax

Overview

sympy2jax

Turn SymPy expressions into parametrized, differentiable, vectorizable, JAX functions.

All SymPy floats become trainable input parameters. SymPy symbols become columns of a passed matrix.

Installation

pip install git+https://github.com/MilesCranmer/sympy2jax.git

Example

import sympy
from sympy import symbols
import jax
import jax.numpy as jnp
from jax import random
from sympy2jax import sympy2jax

Let's create an expression in SymPy:

x, y = symbols('x y')
expression = 1.0 * sympy.cos(x) + 3.2 * y

Let's get the JAX version. We pass the equation, and the symbols required.

f, params = sympy2jax(expression, [x, y])

The order you supply the symbols is the same order you should supply the features when calling the function f (shape [nrows, nfeatures]). In this case, features=2 for x and y. The params in this case will be jnp.array([1.0, 3.2]). You pass these parameters when calling the function, which will let you change them and take gradients.

Let's generate some JAX data to pass:

key = random.PRNGKey(0)
X = random.normal(key, (10, 2))

We can call the function with:

f(X, params)

#> DeviceArray([-2.6080756 ,  0.72633684, -6.7557726 , -0.2963162 ,
#                6.6014843 ,  5.032483  , -0.810931  ,  4.2520013 ,
#                3.5427954 , -2.7479894 ], dtype=float32)

We can take gradients with respect to the parameters for each row with JAX gradient parameters now:

jac_f = jax.jacobian(f, argnums=1)
jac_f(X, params)

#> DeviceArray([[ 0.49364874, -0.9692889 ],
#               [ 0.8283714 , -0.0318858 ],
#               [-0.7447336 , -1.8784496 ],
#               [ 0.70755106, -0.3137085 ],
#               [ 0.944834  ,  1.767703  ],
#               [ 0.51673377,  1.4111717 ],
#               [ 0.87347716, -0.52637756],
#               [ 0.8760679 ,  1.0549792 ],
#               [ 0.9961824 ,  0.79581654],
#               [-0.88465923, -0.5822907 ]], dtype=float32)

We can also JIT-compile our function:

compiled_f = jax.jit(f)
compiled_f(X, params)

#> DeviceArray([-2.6080756 ,  0.72633684, -6.7557726 , -0.2963162 ,
#                6.6014843 ,  5.032483  , -0.810931  ,  4.2520013 ,
#                3.5427954 , -2.7479894 ], dtype=float32)

Turning SymPy expressions into JAX functions

Related tags

Overview

sympy2jax

Installation

Example

Owner

Miles Cranmer

The official implementation of ELSA: Enhanced Local Self-Attention for Vision Transformer

Seach Losses of our paper 'Loss Function Discovery for Object Detection via Convergence-Simulation Driven Search', accepted by ICLR 2021.

WiFi-based Multi-task Sensing

Segmentation models with pretrained backbones. Keras and TensorFlow Keras.

ConformalLayers: A non-linear sequential neural network with associative layers

In real-world applications of machine learning, reliable and safe systems must consider measures of performance beyond standard test set accuracy

Serving PyTorch 1.0 Models as a Web Server in C++

This repository contains the code for Direct Molecular Conformation Generation (DMCG).

Implementation of hyperparameter optimization/tuning methods for machine learning & deep learning models

Using VapourSynth with super resolution models and speeding them up with TensorRT.

RoadMap and preparation material for Machine Learning and Data Science - From beginner to expert.

Embodied Intelligence via Learning and Evolution

Breaking Shortcut: Exploring Fully Convolutional Cycle-Consistency for Video Correspondence Learning

Anomaly detection in multi-agent trajectories: Code for training, evaluation and the OpenAI highway simulation.

Python Library for learning (Structure and Parameter) and inference (Statistical and Causal) in Bayesian Networks.

The official code of "SCROLLS: Standardized CompaRison Over Long Language Sequences".

Simulating Sycamore quantum circuits classically using tensor network algorithm.

Official codebase for Pretrained Transformers as Universal Computation Engines.

Shitty gaze mouse controller

Reliable probability face embeddings