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

U-2-Net: U Square Net - Modified for paired image training of style transfer

Towards Ultra-Resolution Neural Style Transfer via Thumbnail Instance Normalization

A Streamlit component to render ECharts.

My coursework for Machine Learning (2021 Spring) at National Taiwan University (NTU)

pybaum provides tools to work with pytrees which is a concept burrowed from JAX.

Leveraging Social Influence based on Users Activity Centers for Point-of-Interest Recommendation

Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling

HistoKT: Cross Knowledge Transfer in Computational Pathology

Open source implementation of AceNAS: Learning to Rank Ace Neural Architectures with Weak Supervision of Weight Sharing

Retrieval.pytorch - The code we used in [2020 DIGIX]

This repository contains the source codes for the paper AtlasNet V2 - Learning Elementary Structures.

An implementation of the research paper "Retina Blood Vessel Segmentation Using A U-Net Based Convolutional Neural Network"

TextWorld is a sandbox learning environment for the training and evaluation of reinforcement learning (RL) agents on text-based games.

PoseViz – Multi-person, multi-camera 3D human pose visualization tool built using Mayavi.

根据midi文件演奏“风物之诗琴”的脚本 "Windsong Lyre" auto play

This is a Python wrapper for TA-LIB based on Cython instead of SWIG.

A project to make Amazon Echo respond to sign language using your webcam

Stable Neural ODE with Lyapunov-Stable Equilibrium Points for Defending Against Adversarial Attacks

A bunch of random PyTorch models using PyTorch's C++ frontend

Learning recognition/segmentation models without end-to-end training. 40%-60% less GPU memory footprint. Same training time. Better performance.

Turning SymPy expressions into JAX functions

Related tags

Overview

sympy2jax

Installation

Example

Owner

Miles Cranmer

U-2-Net: U Square Net - Modified for paired image training of style transfer

Towards Ultra-Resolution Neural Style Transfer via Thumbnail Instance Normalization

A Streamlit component to render ECharts.

My coursework for Machine Learning (2021 Spring) at National Taiwan University (NTU)

pybaum provides tools to work with pytrees which is a concept burrowed from JAX.

Leveraging Social Influence based on Users Activity Centers for Point-of-Interest Recommendation

Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling

HistoKT: Cross Knowledge Transfer in Computational Pathology

Open source implementation of AceNAS: Learning to Rank Ace Neural Architectures with Weak Supervision of Weight Sharing

Retrieval.pytorch - The code we used in [2020 DIGIX]

This repository contains the source codes for the paper AtlasNet V2 - Learning Elementary Structures.

An implementation of the research paper "Retina Blood Vessel Segmentation Using A U-Net Based Convolutional Neural Network"

​TextWorld is a sandbox learning environment for the training and evaluation of reinforcement learning (RL) agents on text-based games.

PoseViz – Multi-person, multi-camera 3D human pose visualization tool built using Mayavi.

根据midi文件演奏“风物之诗琴”的脚本 "Windsong Lyre" auto play

This is a Python wrapper for TA-LIB based on Cython instead of SWIG.

A project to make Amazon Echo respond to sign language using your webcam

Stable Neural ODE with Lyapunov-Stable Equilibrium Points for Defending Against Adversarial Attacks

A bunch of random PyTorch models using PyTorch's C++ frontend

Learning recognition/segmentation models without end-to-end training. 40%-60% less GPU memory footprint. Same training time. Better performance.

TextWorld is a sandbox learning environment for the training and evaluation of reinforcement learning (RL) agents on text-based games.