Spaces:
Sleeping
Sleeping
File size: 5,627 Bytes
2ceb526 68b3673 2ceb526 68b3673 2ceb526 fb950bb 2ceb526 d0788ef 2ceb526 fb950bb 8c55475 2ceb526 fb950bb 2ceb526 5b978f9 2ceb526 8c55475 55e3b83 2ceb526 44216ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 |
import functools as ft
import sympy
import string
import random
# Special since need to reduce arguments.
MUL = 0
ADD = 1
_jnp_func_lookup = {
sympy.Mul: MUL,
sympy.Add: ADD,
sympy.div: "jnp.div",
sympy.Abs: "jnp.abs",
sympy.sign: "jnp.sign",
# Note: May raise error for ints.
sympy.ceiling: "jnp.ceil",
sympy.floor: "jnp.floor",
sympy.log: "jnp.log",
sympy.exp: "jnp.exp",
sympy.sqrt: "jnp.sqrt",
sympy.cos: "jnp.cos",
sympy.acos: "jnp.acos",
sympy.sin: "jnp.sin",
sympy.asin: "jnp.asin",
sympy.tan: "jnp.tan",
sympy.atan: "jnp.atan",
sympy.atan2: "jnp.atan2",
# Note: Also may give NaN for complex results.
sympy.cosh: "jnp.cosh",
sympy.acosh: "jnp.acosh",
sympy.sinh: "jnp.sinh",
sympy.asinh: "jnp.asinh",
sympy.tanh: "jnp.tanh",
sympy.atanh: "jnp.atanh",
sympy.Pow: "jnp.power",
sympy.re: "jnp.real",
sympy.im: "jnp.imag",
sympy.arg: "jnp.angle",
# Note: May raise error for ints and complexes
sympy.erf: "jsp.erf",
sympy.erfc: "jsp.erfc",
sympy.LessThan: "jnp.less",
sympy.GreaterThan: "jnp.greater",
sympy.And: "jnp.logical_and",
sympy.Or: "jnp.logical_or",
sympy.Not: "jnp.logical_not",
sympy.Max: "jnp.max",
sympy.Min: "jnp.min",
sympy.Mod: "jnp.mod",
}
def sympy2jaxtext(expr, parameters, symbols_in):
if issubclass(expr.func, sympy.Float):
parameters.append(float(expr))
return f"parameters[{len(parameters) - 1}]"
elif issubclass(expr.func, sympy.Integer):
return f"{int(expr)}"
elif issubclass(expr.func, sympy.Symbol):
return f"X[:, {[i for i in range(len(symbols_in)) if symbols_in[i] == expr][0]}]"
else:
_func = _jnp_func_lookup[expr.func]
args = [sympy2jaxtext(arg, parameters, symbols_in) for arg in expr.args]
if _func == MUL:
return ' * '.join(['(' + arg + ')' for arg in args])
elif _func == ADD:
return ' + '.join(['(' + arg + ')' for arg in args])
else:
return f'{_func}({", ".join(args)})'
jax_initialized = False
jax = None
jnp = None
jsp = None
def _initialize_jax():
global jax_initialized
global jax
global jnp
global jsp
if not jax_initialized:
import jax as _jax
from jax import numpy as _jnp
from jax.scipy import special as _jsp
jax = _jax
jnp = _jnp
jsp = _jsp
def sympy2jax(expression, symbols_in, selection=None):
"""Returns a function f and its parameters;
the function takes an input matrix, and a list of arguments:
f(X, parameters)
where the parameters appear in the JAX equation.
# Examples:
Let's create a function in SymPy:
```python
x, y = symbols('x y')
cosx = 1.0 * sympy.cos(x) + 3.2 * y
```
Let's get the JAX version. We pass the equation, and
the symbols required.
```python
f, params = sympy2jax(cosx, [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:
```python
key = random.PRNGKey(0)
X = random.normal(key, (10, 2))
```
We can call the function with:
```python
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:
```python
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:
```python
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)
```
"""
_initialize_jax()
global jax_initialized
global jax
global jnp
global jsp
parameters = []
functional_form_text = sympy2jaxtext(expression, parameters, symbols_in)
hash_string = 'A_' + str(abs(hash(str(expression) + str(symbols_in))))
text = f"def {hash_string}(X, parameters):\n"
if selection is not None:
# Impose the feature selection:
text += f" X = X[:, {list(selection)}]\n"
text += " return "
text += functional_form_text
ldict = {}
exec(text, globals(), ldict)
return ldict[hash_string], jnp.array(parameters)
|