That is still allocating in the inner loop of the ODE solve because of the matmul.
I suggest looking at the following resources:
https://diffeq.sciml.ai/stable/tutorials/faster_ode_example/
That is still allocating in the inner loop of the ODE solve because of the matmul.
I suggest looking at the following resources:
https://diffeq.sciml.ai/stable/tutorials/faster_ode_example/