loss function for Physics informed neural network

[sojohan]

Hi Jax experts

I am trying to build a NN to solve the pendulum ODE as Pinn. Unfortunately, I got the following error.

expected a callable value, got traced<shapedarray(float32[1])

And my loss function looks like this

@jax.jit
def squared_error_dl(t, y):
    #pred_theta = partial(nn_pinn.apply)(params,t)
    #pred_omega=jax.jacfwd(pred_theta,t)(t)
    pred_theta,pred_omega=theta_function(params,t)
    pred=jnp.concatenate((pred_omega,pred_theta))
    return jnp.square(y-pred)
@jax.jit      
def squared_error_eq(t):
    d_omega=jax.jacfwd(pred_omega,t)
    d_eq=-(g/l)*jnp.sin(pred_theta)
    return jnp.square(d_omega-d_eq)

def total_loss(t_batched,y_batched):
    mse_dl = jnp.mean(jax.vmap(squared_error_dl)(t_batched,y_batched)) 
    mse_eq=  jnp.mean(jax.vmap(squared_error_eq)(t_batched))
    return mse_dl + mse_eq

[jakevdp]

It's going to be hard to help here without more information. Can you paste the full traceback, rather than just the one-line error message?