Replace multiple applications of single dispatch with multimethod multiple dispatch

[SamWitty]

Prior to this small refactoring PR, dispatching on both the type of the dynamical system and the type of the solver involved chaining two functools.singledispatchmethod definitions in which the order of arguments (and thus the particular variable being dispatched on) was shuffled around. Instead, this PR has each backend interface method simulate, simulate_trajectory, and get_next_interruptions_dynamic dispatch on both dynamical system and solver type jointly. This substantially reduces indirection, and the development burden when implementing a new backend or dynamical system formalism.

[agrawalraj]

Look great! Could you also re-run the notebook too?

[eb8680]

I sympathize with the goal of this PR, but I'm skeptical about using multimethod this way and a number of the other changes here. Maybe a simpler thing to do would be to stop distinguishing between subtypes of the Dynamics interface (e.g. ODEDynamics), and only dispatch on the solver.

[eb8680]

multimethod dispatches on all of the arguments to _simulate, so what does this registration mean?

[SamWitty]

My understanding is that this registration means that we dispatch only on the first two arguments of _simulate, and not the remaining. This isn't documented in multimethod from what I could find, but when I remove the arguments I get the default NotImplementedError for _simulate or the other dispatches.

[eb8680]

I don't know if that's precisely true (I would guess it implicitly provides a type for the remaining arguments, maybe object or whatever is in the signature?), but I don't think it's a good idea to rely on undocumented behavior in upstream dependencies since it can change without warning in ways that are difficult or impossible to work around.

[SamWitty]

Good point. I'll try and debug why the default dispatch on the full collection of arguments isn't working as expected.

[SamWitty]

@eb8680 , after a bit of digging around I've found two options that rely less on undocumented behavior and pass tests. I have not been able to get multimethod to work with parametric types fully specified.

Option 1: Pass unparameterized types to the *args of register and leave fully parametric types explicit in the function signature. E.g.

@_simulate.register(ODEDynamics, TorchDiffEq, State, torch.Tensor, torch.Tensor)
def torchdiffeq_ode_simulate(
    dynamics: ODEDynamics[torch.Tensor, torch.Tensor],
    solver: TorchDiffEq,
    initial_state: State[torch.Tensor],
    start_time: torch.Tensor,
    end_time: torch.Tensor,
) -> State[torch.Tensor]:
    timespan = torch.stack((start_time, end_time))
    trajectory = _torchdiffeq_ode_simulate_inner(
        dynamics, initial_state, timespan, **solver.odeint_kwargs
    )
    return trajectory[..., -1].to_state()

Option 2: Remove type parameters from function signature and use default register without any *args. E.g.

@_simulate.register
def torchdiffeq_ode_simulate(
    dynamics: ODEDynamics,
    solver: TorchDiffEq,
    initial_state: State,
    start_time: torch.Tensor,
    end_time: torch.Tensor,
) -> State:
    timespan = torch.stack((start_time, end_time))
    trajectory = _torchdiffeq_ode_simulate_inner(
        dynamics, initial_state, timespan, **solver.odeint_kwargs
    )
    return trajectory[..., -1].to_state()

Which of these two would you prefer? Alternatively, would you like something else?

Unrelated: I think there is some misuse of types a bit scattered throughout the module, so a separate type refactoring PR (ideally pair programmed) might be nice before merging into master.

[SamWitty]

I sympathize with the goal of this PR, but I'm skeptical about using multimethod this way and a number of the other changes here. Maybe a simpler thing to do would be to stop distinguishing between subtypes of the Dynamics interface (e.g. ODEDynamics), and only dispatch on the solver.

I actually quite like the idea of dispatching on both the type of the dynamical system and the solver, as this clearly communicates to the user what restrictions we impose/assume about the dynamical system. While we don't impose any explicit restrictions on what's inside the diff method of the ODEDynamics class now, we could, and we probably should soon.

@eb8680 , could you share a bit more about your skepticism? Is it just your comment above about the register decorator, or other things as well?

[eb8680]

• I think dynamical system types like ODEDynamics are not fundamental features of the design, but rather heuristics standing in for restrictions imposed by particular families of solvers, so we shouldn't bake them into the semantics of simulate.
• I prefer having Dynamics be a Protocol rather than a concrete Generic base type, because that constrains our code to have to work entirely in terms of that extremely simple interface.
• I don't think it makes sense to dispatch on all arguments of these operations or enable bigger changes to their signatures like adding or removing arguments.
• I would prefer not to add new mandatory installation dependencies unless we have to. Once Each module should have its own installation requirements #202 is addressed I can be a little more lax about this.
• I'm generally not crazy about multiple dispatch as a design pattern in Python because it's slow and doesn't play nicely with the type system (both in the sense that the runtime dispatch mechanism doesn't work properly with generic types and in the sense that it may throw away generic and parametrized types that could be used by static type checkers or interfere with the use of generic function types elsewhere). There are some circumstances where it does makes sense, but Python isn't Julia and it should be used with care.

[SamWitty]

• I think dynamical system types like ODEDynamics are not fundamental features of the design, but rather heuristics standing in for restrictions imposed by particular families of solvers, so we shouldn't bake them into the semantics of simulate.

I generally agree with this point. Does this mean that onus is on the user to know what the restrictions are on any particular solver? Alternatively, are there static or runtime checks that each solver can implement that enforce their restrictions? I'm very nervous about a user e.g. adding pyro.sample statements in the body of a diff method and then using the TorchDiffEq solver.

• I prefer having Dynamics be a Protocol rather than a concrete Generic base type, because that constrains our code to have to work entirely in terms of that extremely simple interface.

I see. I think I need to read up on the distinction here.

• I don't think it makes sense to dispatch on all arguments of these operations or enable bigger changes to their signatures like adding or removing arguments.

This could be addressed by passing in object explicitly in the register decorator.

• I would prefer not to add new mandatory installation dependencies unless we have to. Once Each module should have its own installation requirements #202 is addressed I can be a little more lax about this.

Fair. It is a bit sad to have this be the only other dependency besides Pyro, and it's not obvious how stable it is.

• I'm generally not crazy about multiple dispatch as a design pattern in Python because it's slow and doesn't play nicely with the type system (both in the sense that the runtime dispatch mechanism doesn't work properly with generic types and in the sense that it may throw away generic and parametrized types that could be used by static type checkers or interfere with the use of generic function types elsewhere). There are some circumstances where it does makes sense, but Python isn't Julia and it should be used with care.

Also fair. I've experienced this a bit already (see above workaround to your question about register.

[eb8680]

Does this mean that onus is on the user to know what the restrictions are on any particular solver?

Keep in mind that this is currently the case - ODEDynamics doesn't actually enforce any restrictions on the contents of diff, they're just implied by the name.

are there static or runtime checks that each solver can implement that enforce their restrictions? I'm very nervous about a user e.g. adding pyro.sample statements in the body of a diff method and then using the TorchDiffEq solver.

We could implement simple validation effect handlers that get applied to diff and raise errors at runtime when certain operations are encountered. A more powerful type system would let us enforce those restrictions statically, but I don't think that's currently doable with mypy and Pyro's handlers' janky string-based dynamic dispatch mechanism.

[eb8680]

This could be addressed by passing in object explicitly in the register decorator.

Hmm, no, my point was that someone could write a signature with a finer type than object for those other arguments like timespan, which is a form of extensibility I don't think we should support.

[SamWitty]

Closing this PR in favor of an alternative that removes ODEDynamics. Currently a WIP on my local machine dealing with circular imports, but will create a PR soon.