First DDE version

[thibmonsel]

New files :

• discontinuity.py that does the root finding during integration steps.

Modified files:

• integrate.py changed a bit of the code but essentially looks the same but with more if statements. There is also the discontinuity handling before each integration step done. Added 2 new arguments to _State (discontinuities, discontinuities_save_index)
• constant.py does the discontinuity checking and returns the next integration step. But as said in WIP the prevbefore and nextafter are done in the loop()

Followed your suggestion regarding dropping y0_history and putting it in y0. However by doing this we must pass y0 to the loop function now. Haven't done the PyTree handling of delays yet. Only works for constant stepsize controller, doing adaptive now.

PS : I dont have the save saving format as you so terms.py shows some deletion and addition for no reason ....

Boilerplate code for a dde :

def vector_field(t, y, args, *, history):
        return 1.8 * y * (1 - history[0])

    delays = [lambda t, y, args: 1]
    y0_history = lambda t: 1.2
    discontinuity = (0.0,)

    made_jump = discontinuity is None
    t0, t1 = 0.0, 100.0
    ts = jnp.linspace(t0, t1, 1000)
    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(vector_field),
        diffrax.Dopri8(),
        t0=ts[0],
        t1=ts[-1],
        dt0=ts[1] - ts[0],
        y0=y0_history,
        max_steps=2 ** 16,
        stepsize_controller=diffrax.ConstantStepSize(),
        saveat=diffrax.SaveAt(ts=ts, dense=True),
        delays=delays,
        discontinuity=discontinuity,
        made_jump=made_jump,
    )

    plt.plot(sol.ts, sol.ys)
    plt.show()

[patrick-kidger]

Okay, so there's a lot of spurious changes here, mostly due to unneeded formatting changes. Take a look at CONTRIBUTING.md, and in particular the pre-commit hooks. These will autoformat etc. the code. I'll be able to do a proper review then.

Regarding passing y0_history into loop: I'm thinking what we should do is something like:

def diffeqsolve(y0, delays, ...):
    if delays is None:
        y0_history = None
    else:
        y0_history = y0
        y0 = y0_history(t0)

    adjoint.loop(..., y0_history=y0_history)

so that internally we still disambiguate between y0 and y0_history.

Regarding the changes to constant step sizing: hmm, this seems strange to me. I don't think we should need to change any stepsize controller at all. I think the stepsize controller changes we need to make (due to discontinuities) should happen entirely within integrate.py, so that they can apply to all stepsize controllers. (e.g. even user-specified stepsize controllers, that don't have any special support)

[thibmonsel]

Hello,
Im back with some updates !

1/ I used the pre-commit hooks but still have one small spurious change in terms.py 😲 .
1/ y0 and y0_history are disambiguated !
2/ The controllers are untouched and better for modular code.
3/ Discontinuity handling is done in loop file.

4/ 1 edges cases was found, i haven't thought too much and did a "sloppy" fix for now (https://github.com/patrick-kidger/diffrax/pull/169/files#r991488410). Essentially it comes when we integrate a step from tprev to tnext and the integration bound new_tprev=tnext is right next to a discontinuity forcing to redo the step and having an new interval being [tnext ; tnext + epsilon = discontinuity_jump]. This in hand with the code handling will put tprev > tnext and throw an error.
Regarless, for most cases the solver seems to work

Things not done :
1/ y0_history is still not a PyTree of callable but when that will be done i suppose to change its structure one should use eqx.tree_at ?

[thibmonsel]

Latest commit does what you suggested

• Reverted back to not touching the state.tnext and such.
• Integrated your bissection code into the discontinuity detection.
• Only search for discontinuities if step is rejected.
• delays are PyTrees

To do :

• code cleanup ie removing discontinuity.py that now is obsolete (will do next commit when something stable is available)
• implicit stepping when the state.tnext - state.tprev > min(delays)
  • regarding this I checked the julia paper (4.2), I understand the issue there but the explanation on how to handle it is rather opaque i find... (*). In this case, the use of the continuous extension (ie dense_interp from dense_infos) makes the overall method implicit even if the discrete method we are using is explicit (this is called overlapping apparently).

Comments :

• Regarding the first commit where lots of thing where changed, I wasn't aware of the 2 point of views (philosophies) on how to integrate DDEs. The first is to track the discontinuities and the other relies on error estimation of the solver's method. On the latter discontinuity tracking is somewhat useless. Julia uses therefore the second idea.
• Since the implicit stepping (cf *) is not done some integrations are faulty.

[thibmonsel]

Latest commit has :

• added a discont_update in the discontinuity checking part to make sure we correctly update state.discontinuity
• added a discont_check argument in diffeqsolve to check for breaking points at each step
• implemented the implicit step in history_extrapolation_implicit
• removed the discontinuity.py so everything is in integrate.py
• enhanced discontinuity tracking with another functionality to check also for roots in subintervals of the [tprev, tnext], this helps to not miss any potential jumps that would otherwise be undetected. Not sure, but i think the vmaping here doesn't destroy the computational gain of the unvmap ?

[patrick-kidger]

Okay I've reviewed about half of it! I'll try and get to the other half in the next week or so.

Overall I like where this is going. I think we now have some first implementations for every piece.

Regarding my comments about avoiding in-place updates. I think this is probably doable by updating HistoryVectorField to operate on three regions: y0_history, the recorded dense_infos, and finally the current step. (Much like how it is operating over two regions at the moment.) I think this should allow us to generate efficient code.

[thibmonsel]

Latest commit updates some remarks of code comment after latest review #169 (review).

I've bundled together as you said the delays term together for an easier API.

class _Delays(eqx.Module):
    delays: Optional[PyTree[Callable]]
    initial_discontinuities: Union[Array, Tuple]
    max_discontinuities: Int
    recurrent_checking: Bool
    rtol: float
    atol: float
    eps : float

delays is the regular delayed term definition, initial_discontinuities corresponds to the discontinuities that the user needs to give to get proper DDE integration (I refer you to my response #169 (comment)) , discont_checking is now recurrent_checking and the atol, rtolare for the Newton solver in the implicit step and eps the tolerance error or the step.

What you suggested is great for later iterations because we can just slap any new arguments into _Delays !

What needs to be done :

• First DDE version  #169 (comment) Didn't transfer code snipped because I wasn't sure which one it was.
• First DDE version  #169 (comment) Finding a way to not redo the solver.step when we do implicit steps.
• Take a look into the coding of dense_infos and dense_ts and in-place updates

[patrick-kidger]

Okay, let's focus on this bit before we move on to discussing the discontinuity handling below. I'll leave you to make the changes already discussed here. Let me know if any of them aren't clear.

[thibmonsel]

Okay, let's focus on this bit before we move on to discussing the discontinuity handling below. I'll leave you to make the changes already discussed here. Let me know if any of them aren't clear.

Sounds good, I'll take care of the first bullet point later, the second one should be ok on my side however i'd like to have your take on the third one with the in-place operations (for dense_ts, dense_infos) since this is some very sharp JAX bit 🔪 .

[patrick-kidger]

Sure thing. I'm suggesting that _HistoryVectorField should look something like this:

class _HistoryVectorField(eqx.Module):
    ...
    tprev: float
    tnext: float
    dense_info: PyTree[Array]
    interpolation_cls: Type[AbstractLocalInterpolation]

    def __call__(self, t, y, args):
        ...
        if self.dense_interp is None:
            ...
        else:
            for delay in delays:
            
                delay_val = delay(t, y, args)
                alpha_val = t - delay_val

                is_before_t0 = alpha_val < self.t0
                is_before_tprev = alpha_val < self.tprev
                at_most_t0 = jnp.min(alpha_val, self.t0)
                t0_to_tprev = jnp.clip(alpha_val, self.t0, self.tprev)
                at_least_tprev = jnp.max(alpha_val, self.tprev)

                step_interpolation = self.interpolation_cls(
                    t0=self.tprev, t1=self.tnext, **self.dense_info
                )
                switch = jnp.where(is_before_t0, 0, jnp.where(is_before_tprev, 1, 2))
                history_val = lax.switch(switch, [lambda: self.y0_history(at_most_t0),
                                                  lambda: self.dense_interp(t0_to_tprev),
                                                  lambda: step_interpolation.evaluate(at_least_tprev)])
                history_vals.append(history_val)
        ...
        return ...

And then when it is called inside the implicit routine:

def body_fun(val):
    dense_info, ... = val
    ...
    _HistoryVectorField(..., state.tprev, state.tnext, dense_info, solver.interpolation_cls)
    ...
    return new_dense_info, ...

[thibmonsel]

Latest commit should have handle all of the issues mentionned above.

• Not do an extra step after doing the implicit step by integrating de-facto the explicit step in the lax.for_loopof the function history_extrapolation_implicit with the conditional your proposed (First DDE version  #169 (comment))
• In-place update handled with _HistoryVectorField

Discussion/Bottleneck for implicit step

Regarding the implicit step we have a issue when it comes to large steps because an step_interpolation with only 2 points won't suffice. This depends on the fact that the snippet below is indeed a 2 point interpolation :

step_interpolation = self.interpolation_cls(t0=self.tprev, t1=self.tnext, **self.dense_info)

To elaborate a bit more, if we have an implicit step from state.tprev to state.tnext. Our associated history function for the equation y'(t) = f(t, y(t-tau)) will be known from state.tprev up to state.tprev + tau and from state.prev+tau to state.tnext we need its extrapolation. If the history function in the interval [state.prev+tau : state.tnext] is non monotonous (we could image half period of a sinus for example) well a 2 point extrapolation won't capture the function correctly but we would need rather 10 points lets say to get a good estimate. To this regard we would need also to change the conditioning of your implict step from

_pred = (((y - y_prev) / y) > delays.eps).any()

to something that checks the MSE of the extrapolated history function before and after the integration step. Not sure with this in mind the _HistoryVectorField from #169 (comment) as is will do the trick.

This also impacts too the population of the ys in dense_ts since the values are interpolated with the computed steps of y. (If we go with an implicit step from state.tprev to state.tnext and we have to save some points in between, the current procedure is to use a 2 point (yprev=state.y and ynext=y) interpolation right ?) With that being said, this discussion is only relevant if the time mesh that we have (ie dense_ts) is precise enough.

[patrick-kidger]

Our associated history function for the equation y'(t) = f(t, y(t-tau)) will be known from state.tprev up to state.tprev + tau

I don't think this is true. Anything after tprev hasn't been evaluated yet. The whole [tprev, tnext] region is initialised as an extrapolation from the previous step.

Regarding 2-point interpolations: this isn't the case for most solvers. Each solver evaluates various intermediate quantities during its step (e.g. the stages of an RK solver) and these also feed into the interpolation.

Even if were, though, I don't think it matters: we just need to converge to a solution of the implicit problem.

[thibmonsel]

Ok, this makes sense, so i'll take back what I said in my bottleneck "Discussion/Bottleneck for implicit step", thanks for the clarification ! The _pred condition should be on the dense_info then !

Edit : _pred condition as it seems to be working fine !
For this part of the code, i'd say its ready for a review before moving to discontinuity checking ! (unconstrained time stepping works well)

[thibmonsel]

Relevant changes are made :

• All the code is now managed with PyTree operations (i think)
• Moved Delays and implicit step into a new file (but there is a circular import there ... not sure if this is an issue)
• Removed eps features of the implicit stepping

In order to get something backprop compatible in history_extrapolation_implicit I'll need to use your bounded_while_loop instead of lax.while_loopin the meantime until your PR is merged ?

[patrick-kidger]

Okay, this is probably the final round of review on the implicit stepping. (Nice work!) I'll go over the discontinuity checking at some point shortly.

Regarding backprop through the lax.while_loop, the correct thing to do here is actually to use the implicit function theorem, e.g. as is already done with Newton's method. This can be wired up using the misc.implicit_jvp helper. Let me know if you're not familiar with this and I can give you some pointers on how this works.

[thibmonsel]

Great news !
Im not familiar with the wiring of the implicit_jvp up for some pointers.
Other then that relevant changes were made

[patrick-kidger]

Nit: I think we only want to expose Delays in the public interface. Everything else is Diffrax-internal.

[thibmonsel]

I'm sorry not sure to understand what you mean

[patrick-kidger]

What if we hit max_steps here? I think we need a mechanism in this whole discontinuity procedure to allow the discontinuity-finding to fail. (Namely, reject the step and use some more naive way of picking tnext, e.g. tprev plus 0.1 times the current interval length or something.)

[thibmonsel]

If we hit max_steps that would just means that the NewtonSolver would take the lead with the current root approximation ? Speaking from experience, since each sub_tprev and sub_tnext used are pretty small intervals this would probably not happen ? The interval [tprev, tnext] is splitted into N sub intervals where root tracking is done.

[patrick-kidger]

I don't think we can ever hit this branch?
(And indeed, as currently written it would crash, since it doesn't also return a bool.)

[thibmonsel]

Yup we never would hit this, no problem should be encountered since NewtonNonlinearSolver return NaNs if it fails right ? (if that is the case we should be good then )

[patrick-kidger]

Alright, on to the next block of code!

As for implicit_jvp -- I actually have some in-progress work that may simplify this. If you'd like to be able to backpropagate through this code soon-ish then I can expand on what I mean here? (But if it's not a rush then I'm happy to put a pin in that for now.)

[thibmonsel]

• Relevant changes were made for the last code block, everything is in delays.py now !
• Added 2 new attributes for Delays : nb_sub_intervals and max_steps
• Lots of wrapper and deletion makes the code more readable !

As for implicit_jvp, backpropagating throught this part would be great !
Looking into this #169 (comment) a bit deeper now

[patrick-kidger]

This is looking good. Let's start adding this into the documentation and so on:

• Add a new doc page under the "Advanced API" section.
• Add a short example. (See the steady-state example for inspiration.)

Let's also add a test_delays.py file, checking all the edge cases. Off the top of my head:

• Basic checks that the solver actually runs without crashing.
• Numerical checks that the expected solution is obtained.
• Check what happens when we exceed Delays.max_discontinuities.
• Check what happens when we exceed Delays.max_steps.
• Test combining delays with stochastic terms.
• Test combining delays with PIDController(jump_ts=...).
• Test a 'smooth' DDE (whose initial history is such that there is no discontinuity), and check that this can be solved without any expensive discontinuity detection at all. (Actually, I don't think this is possible at the moment -- maybe the recurrent_checking argument should be generalised: discontinuity_checking=True/False/None for always check / only check on rejected steps / never check?)
• Test DDEs that hit both the implicit-step and the explicit-step branches, and that those branches are taken. (To test this: perhaps we can count the number of implicit and explicit steps, and return that in the auxiliary stats.)

Other changes that come to mind:

• Another good auxiliary statistic could be how many discontinuities were encountered. (Including any discontinuities arising from jump_ts?)
• I think the numerics might have still have a subtle bug from the lack of something like

  diffrax/diffrax/integrate.py

  Line 82 in 5f5a121

  def _clip_to_end(tprev, tnext, t1, keep_step):

  not being used every time we clip tnext.

---

As for implicit_jvp. Fix some function f, and define y(θ) as being the value of y satisfying f(y, θ) = 0. (And we assume there is a unique such y.) Then we can see that the function f has implicitly defined a function θ -> y(θ).

We can now seek to evaluate the derivative dy/dθ. This actually involves a linear solve, and is what implicit_jvp does. This is pretty easy to do. Probably the best reference is:
http://implicit-layers-tutorial.org/implicit_functions/
Also see the final equation of this section:
https://en.wikipedia.org/wiki/Implicit_function_theorem#Statement_of_the_theorem

You can see this in action for the existing nonlinear solvers, which use the implicit function theorem to differentiate through their root-finding operation.

[patrick-kidger]

Before merging, we'll have to rebase your branch onto the current main version of Diffrax. When that happens, heads-up that these imports will change to:

import equinox.internal as eqxi
from equinox.internal import ω

eqxi.unvmap_any

[thibmonsel]

Thanks for the review Patrick :) , so from what I understood what i need to do is to create a new class DDEImplicitNonLinearSolve (for e.g) that inherits from AbstractNonlinearSolver in order to use implicit_jvp. Basically i'll need to rewrite the _solve method and then in history_extrapolation_implicit call in an instance of DDEImplicitNonLinearSolve.

class DDEImplicitNonLinearSolve(AbstractNonlinearSolver):
    def _solve(
        self,
        fn: callable, # would be terms
        x : Pytree, # here would be state.y
        nondiff_args: PyTree, # here would be all the other args from current history_extrapolation_implicit(...)
        diff_args: PyTree,
    ) 

def history_extrapolation_implicit(...):
     nonlinearsolver = DDEImplicitNonLinearSolve(...)
     results = nonlinearsolver(terms, y, args).root 
     y, y_error, dense_info, solver_state, solver_result = results 
     return y, y_error, dense_info, solver_state, solver_result

If thats the case could you explain how you usually work with your non_diff_args and diff_args since I think its more native to JAX's jvp/vjp and are later on used with implicit_jvp and _rewrite, etc... ?

[patrick-kidger]

Okay, apologies for how long this has taken me to return to! The new Diffrax release is now out, and I should have a bit more time to focus on other larger-scale changes to Diffrax like this.

I've had another look over the code, and for the most part I think this is pretty much ready! So:

• Can you rebase on to the new Diffrax release? (Squashing your commits together may make this easier.) I'll do a check-over after that to make sure things look like they've happened cleanly.
• Can you add a non-neural DDE example as well? Just something short and simple that demonstrates the API, like the existing "Stiff ODE" example.
• The Julia file appears to be entirely commented out -- what's going on with that?
• I think you need to enable the pre-commit hooks; some of your new changes aren't properly formatted etc.

[thibmonsel]

Hi @patrick-kidger , I'm having another go on the DDEs and with all of the changes on Diffrax and Equinox.
I have rebased the code and with it, there is new inner_while_loop/outer_while_loop in loop !

DenseInterpolation (in integrate.py l.238) is inside inner_while_loop/outer_while_loop and all the arguments of DenseInterpolation gets transformed before hand with Equinox's _Buffer object/module and the DenseInterpolation.__post_init__ check fails. I was wondering if there's anyway around it ?

Thanks !

[patrick-kidger]

Ah, I'm guessing it's the ts=state.dense_ts and infos=state.dense_infos that are problematic. Try something like:

ts = state.dense_ts[...]
infos=jtu.tree_map(lambda _, x: x[...], dense_info, state.dense_infos)

to unwrap them from buffers.

Note that unwrapping like this is an unsafe thing to do in general! You have to make the promise to only read from locations that have already been written to with an in-place update -- otherwise you'll get incorrect gradients. (I would definitely add some tests that finite differences match the results from jax.grad, just to be sure. FYI we have a finite_difference_jvp here that you could copy-paste.)

[thibmonsel]

Adding [...] did the trick for ts

ts = state.dense_ts[...]

Unfortunately unwrapping the Buffer with other structure's like DenseInfos didn't seem to work.

infos=jtu.tree_map(lambda x: x[...], state.dense_infos)

The problem came from the _pred argument of the _Buffer

However, I found another way

unwrapped_buffer = jtu.tree_leaves(
    eqx.filter(state.dense_infos, eqx.is_inexact_array),
    is_leaf=eqx.is_inexact_array,
)
unwrapped_dense_infos = dict(zip(state.dense_infos.keys(), unwrapped_buffer))

[patrick-kidger]

Notice how I pass in an additional argument in my previous snippet: jtu.tree_map(lambda _, x: x[...], dense_info, state.dense_infos).

This is because each _Buffer object is actually a PyTree, that's sort of masquerading as an array. So to handle this we actually iterate over a different pytree, that happens to have the correct structure.

The use of buffers is a pretty advanced/annoying detail. I'm pondering using something like Quax to create a safer API for this, but that's a long way down the to-do list.

[thibmonsel]

Indeed your right, I just realised that in the documentation yesterday, this makes sense now !
However, we only have access to dense_info after creating DenseInterpolation.

i.e.

 dense_interp = DenseInterpolation(
                ts=state.dense_ts[...],
                infos = jtu.tree_map(lambda _, x: x[...], dense_info, state.dense_infos), 
                ...
            )

(
    y,
    y_error,
    dense_info,
    solver_state,
    solver_result,
) = history_extrapolation_implicit(
    ...
)

So I would agree that this works if dense_info was available which seems to be not the case. Hence not sure how to do it. Moreover, I did some preliminary testing and doing this #169 (comment) yields wrong gradients. Nonetheless, the code makes the promise to only read from locations that have already been written to with an in-place update so my unwrapping method is/seems erroneous.

[patrick-kidger]

Since I'm starting to see activity here again -- you can track my progress updating Diffrax in #217 and https://github.com/patrick-kidger/diffrax/tree/big-refactor.

(This currently depends on the unreleased versions of Equinox and jaxtyping.)

Mostly there now!