why when i have a np.power in my function jax.grad can't give me the derivitives?

user_6663733 0
Viewed 630

I want to train a simple linear model. these below x and y are my data.

import numpy as np
x = np.linspace(0,1,100)
y = 2 * x + 3 + np.random.randn(100)

f is a function that calculates mean square error over all data.

def f(params, x, y):
  return np.mean(np.power((params['w'] * x + params['b'])-y , 2))


from jax import grad
df = grad(f)
params = dict()
#initialize parameters
params['w'] = 2.4
params['b'] = 10.
df(params, x, y) # I will do this in a loop (implementing gradient decent part

this gives me an error:

FilteredStackTrace: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ConcreteArray

when I clear np.power code works. why?

jax linear-regression numpy python
1 Answers

JAX cannot compute gradients of numpy functions, but it can compute gradients of jax.numpy functions. If you rewrite your code in terms of jax.numpy, it should work for you:

import numpy as np
x = np.linspace(0,1,100)
y = 2 * x + 3 + np.random.randn(100)

import jax.numpy as jnp
def f(params, x, y):
  return jnp.mean(jnp.power((params['w'] * x + params['b'])-y , 2))

from jax import grad
df = grad(f)
params = dict()

params['w'] = 2.4
params['b'] = 10.
df(params, x, y)
# {'b': DeviceArray(14.661432, dtype=float32),
#  'w': DeviceArray(7.3792152, dtype=float32)}

You can read more details in the TracerArrayConversionError documentation page.

user_29378310
Related