diffrax icon indicating copy to clipboard operation
diffrax copied to clipboard

Warning for solving example ODE from docu

Open MaAl13 opened this issue 8 months ago • 3 comments

The following example code from the docu produces a warning. How can i get rid of it?

import jax.numpy as jnp
from diffrax import diffeqsolve, ODETerm, SaveAt, Tsit5


def vector_field(t, y, args):
    prey, predator = y
    α, β, γ, δ = args
    d_prey = α * prey - β * prey * predator
    d_predator = -γ * predator + δ * prey * predator
    d_y = d_prey, d_predator
    return d_y


term = ODETerm(vector_field)
solver = Tsit5()
t0 = 0
t1 = 140
dt0 = 0.1
y0 = (10.0, 10.0)
args = (0.1, 0.02, 0.4, 0.02)
saveat = SaveAt(ts=jnp.linspace(t0, t1, 1000))
sol = diffeqsolve(term, solver, t0, t1, dt0, y0, args=args, saveat=saveat)

Warning:

'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>. Attempting to hash a tracer will lead to an error in a future JAX release. warnings.warn(

MaAl13 avatar Jun 20 '24 14:06 MaAl13