[ ] Implement a Parameters class in the evaluator module that can be passed to JAX and SUNODE ODE systems. Since SUNODE expects a class instance / dict anyway and JAX accepts PyTree objects, this should work. Implementing this also gives some more freedom of providing nested parameter structures for hierarchical models.
This is an optional improvement
Parametersclass in theevaluatormodule that can be passed toJAXandSUNODEODE systems. SinceSUNODEexpects a class instance / dict anyway andJAXacceptsPyTreeobjects, this should work. Implementing this also gives some more freedom of providing nested parameter structures for hierarchical models.