smsharma

smsharma

smsharma

smsharma

import time

import jax
import jax.numpy as jnp

# Global flag to set a specific platform, must be used at startup.
jax.config.update("jax_platform_name", "gpu")

# from diffrax import diffeqsolve, ODETerm, Dopri5, PIDController, SaveAt
from jax.experimental.ode import odeint

smsharma

smsharma

import matplotlib.pylab as pylab
from cycler import cycler
import seaborn as sns

# params = {'text.usetex': True,
#           'font.family': 'serif',
#           'font.serif': ['CMU Serif'],
#           'font.sans-serif': ['CMU Sans Serif'],
#           'legend.fontsize': 17,
#           'legend.fancybox': True,

smsharma

smsharma

smsharma

smsharma