import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np
def softplus(x):
    return jnp.log(1 + jnp.exp(x))
def init_params():
    params = jax.random.normal(key, shape=(241,))
    return params
def linear_model(params, x):
    w0 = params[:80]
    b0 = params[80:160]
    w1 = params[160:240]
    b1 = params[240]
    h = softplus(x*w0 + b0)
    o = jnp.sum(h*w1) + b1
    return o
def loss(derivative, initial_condition, params, model, x):
    dfdx = jax.grad(model, 1)
    dfdx_vect = jax.vmap(dfdx, (None, 0))
    model_vect = jax.vmap(model, (None, 0))
    eq_difference = dfdx_vect(params, x) - derivative(x, model(params, x))
    condition_difference = model(params, 0) - initial_condition
    return jnp.mean(eq_difference ** 2 - condition_difference ** 2)
def dfdx(x, y):
    return -2. * x * y
key = jax.random.PRNGKey(0)
inputs = np.linspace(0, 1, num=401)
params = init_params()
epochs = 2000
learning_rate = 0.0005
# Training Neural Network
for epoch in tqdm(range(epochs)):
    grad_loss = jax.grad(loss)
    gradient = grad_loss(dfdx, 1., params, linear_model, inputs)
    params -= learning_rate*gradient
model_vect = jax.vmap(linear_model, (None, 0))
preds = model_vect(params, inputs)
plt.plot(inputs, jnp.exp(inputs**2), label='exact')
plt.plot(inputs, model_vect(params, inputs), label='approx')
plt.legend()
plt.show()

TypeError: Argument '<function dfdx at 0x7fce88340af0>' of type <class 'function'> is not a valid JAX type.

def loss(params, derivative, initial_condition, model, x):
    dfdx = jax.grad(model, 1)
    dfdx_vect = jax.vmap(dfdx, (None, 0))
    model_vect = jax.vmap(model, (None, 0))
    eq_difference = dfdx_vect(params, x) - derivative(x, model(params, x))
    condition_difference = model(params, 0) - initial_condition
    return jnp.mean(eq_difference ** 2 - condition_difference ** 2)