jax gradiant problem

[ak24watch]

why is it happening and how to resolve

for epoch in range(20000):

def loss_fn(model):
    # compute data loss
    yh = jax.vmap(model)(t_data)
    loss1 = jnp.mean((yh - x_t_data) ** 2)
    # compute differential equations

    def compute_x_t(t):
        y = model(t)  # y has shape (1,)
        return y[0]   # Extract the scalar value
    
    # Get the first derivative function
    dx_dt_fn = jax.grad(compute_x_t)  # this line ok
    
    
    # Get the second derivative function
    d2x_dt_fn = jax.grad(dx_dt_fn)         ==> this line problem
    
    # Apply to all points in t_physics
    dx_t = jax.vmap(dx_dt_fn)(t_physics)   =====> this line ok
    d2x_t = jax.vmap(d2x_dt_fn)(t_physics)     =====> ==> this line problem


  
    physocs = d2x_t + mu * dx_t + k * jax.vmap(model)(t_physics)
    loss2 = jnp.mean(physocs ** 2)*(1e-4)

    # combine the losses
    loss = loss1 + loss2
    return loss
loss, grads = eqx.filter_value_and_grad(loss_fn)(model)
# update
updates, opt_state = optim.update(grads, opt_state, eqx.filter(model, eqx.is_array))
model = eqx.apply_updates(model, updates)
if (epoch + 1) % 1000 == 0:
    print(f"Epoch {epoch + 1}, Loss: {loss:.4f}")
    # plot the results
    yh = jax.vmap(model)(t)
    plot_result(t, x_t, t_data, x_t_data, yh, t_physics, i=epoch)
    plt.show()

[jakevdp]

Hello - please use code blocks to properly format code, and please avoid screenshots of text: instead copy and paste the text itself.

Also, it looks like you posted the same question at #28110. I'm going to close this issue to avoid duplication.