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

Neighbor2Seq: Deep Learning on Massive Graphs by Transforming Neighbors to Sequences

Breast Cancer Classification Model is applied on a different dataset

Kaggleship: Kaggle Notebooks

Code for Efficient Visual Pretraining with Contrastive Detection

An Unsupervised Detection Framework for Chinese Jargons in the Darknet

A generator of point clouds dataset for PyPipes.

Implementation of the Chamfer Distance as a module for pyTorch

Deepfake Scanner by Deepware.

Predicting future trajectories of people in cameras of novel scenarios and views.

A python script to dump all the challenges locally of a CTFd-based Capture the Flag.

Code for our ACL 2021 paper - ConSERT: A Contrastive Framework for Self-Supervised Sentence Representation Transfer

Implementation of QuickDraw - an online game developed by Google, combined with AirGesture - a simple gesture recognition application

Self-Adaptable Point Processes with Nonparametric Time Decays

WORD: Revisiting Organs Segmentation in the Whole Abdominal Region

Unsupervised 3D Human Mesh Recovery from Noisy Point Clouds

Tensorflow implementation of our method: "Triangle Graph Interest Network for Click-through Rate Prediction".

TensorFlow implementation of "TokenLearner: What Can 8 Learned Tokens Do for Images and Videos?"

Converting CPT to bert form for use

DeepI2I: Enabling Deep Hierarchical Image-to-Image Translation by Transferring from GANs

Program your own vulkan.gpuinfo.org query in Python. Used to determine baseline hardware for WebGPU.