Replace a reference to jax.tree_map with tree_util.tree_map

pnkraemer / probdiffeq

Probabilistic solvers for differential equations in JAX. Adaptive ODE solvers with calibration, state-space model factorisations, and custom information operators. Compatible with the broader JAX scientific computing ecosystem.

https://pnkraemer.github.io/probdiffeq/

MIT License

30 stars 2 forks source link

Replace a reference to jax.tree_map with tree_util.tree_map #700

Closed pnkraemer closed 9 months ago