ReflectOn: pick which method's body to rewrite

[shashi]

We need this in ForwardDiff2 to make Dual numbers work on functions which are restricted to concrete number types, for example, f(x::Float64) = x+1.

This will allow us to take the method f(::Float64) and rewrite it while still using a Dual container type.

[shashi]

But this fails when use closures.... I guess it would be nice to pass in the closure as an extra argument and set it to be the #self# somehow.

julia> g(x) = y->x+1                                                                                                                                                                            g (generic function with 1 method)                                                                                                                                                              
julia> Cassette.overdub(Foo(), Cassette.ReflectOn{Tuple{typeof(g(0)), Float64}}(), 9)
ERROR: type ReflectOn has no field x

[vchuravy]

Thanks Shashi, IIUC this is exactly the use-case for Cassette's metadata system, so I am interested in why it doesn't work for you, if it can't fulfill your use-case it might be better to rip it out and not have the complexity of maintaining it.

[oxinabox]

I don't think this is the usecase for the metadata system.
Metadata just carries around extra state, which you might want to record various things.
Its pretty useful for that, except that it sometimes triggers large overhead.

This is about foo(::Dual) is called but doesn't exist,
foo(::Float64) exists,
using overdubbing, one can programatically determine what foo(::Dual) should do, if one has access to foo(::Float64).

Isn't it?

[MikeInnes]

Pretending for the purpose of program semantics (dispatch) that you have only a Float64 while actually storing extra stuff (e.g. perturbations) is an exact description of the metadata system, and this is in fact the exact use case it was designed for. Going down that path opens a huge can of worms (the main reason FD2 isn't using it already) so you want to be wary if you're doing that.

In particular this seems concerning if you call things like constructors with incorrect type information, e.g. Complex{Float64}(Dual(...), ...). This seems like a recipe for disaster, without bringing in the full metadata machinery, but perhaps there's a clever way around it.

[MikeInnes]

[Metadata and Dual containers aren't the only two options, FWIW; a third option is to transform the code so that all values are effectively Duals, including structs, functions etc., and then the constructor issue above is automatically solved. This approach also prevents any kind of perturbation confusion by default. This goes a bit outside of what is easy to support with Cassette's abstractions though.]

[oxinabox]

Pretending for the purpose of program semantics (dispatch) that you have only a Float64 while actually storing extra stuff (e.g. perturbations) is an exact description of the metadata system, and this is in fact the exact use case it was designed for.

It means your extra data has to live in the metadata at basically all times, rather than living in the arguments where it naturally would be.

[MikeInnes]

That's kind of a user interface problem rather than anything fundamental about the metadata system, though. If it's more efficient to pass metadata as arguments the system can and should do that.

[MasonProtter]

I think for API consistency, it should be ReflectOn{f, Tuple{ArgTypes...}} instead of ReflectOn{Tuple{f, ArgTypes...}}

[shashi]

fair enough.

[shashi]

Cassette's tagging system is super slow for now. It quickly breaks inference and results in run-away allocations. I'd be the first one to use it when LICM or whatever is needed makes it faster. That's why for now we're trying to do it with a Dual container type.

This change is necessary to have correctness for differentiating functions like:

f(x::Int64) = x+1
f(x) = x^2

D(f)(1)

FD2 will not work for structs which have strict typing. We've decided to live with that, and that kind of code is not common in Julia.

Complex{Float64}(Dual(...), ...)

That will have to be dealt with by not using ReflectOn in cases where there are no leaf nodes with a valid frule.

Either way, this is a powerful feature which already exists, and since this package is all about the power, we should expose it.

[MikeInnes]

I appreciate the motivation 100%, I'm just not totally clear that there's a good middle ground between respecting Julia semantics fully (limiting both dispatch and type support) and going full-on with the metadata system. Unless you're really quite certain that this doesn't lead you down the rabbit hole (e.g. by showing that using it in FD2 works well in practice) I think Valentin's concern about general usefulness vs maintenance is warranted.

My instinct is that this if you do it in halves, you'll end up with something that's fragile, has a lot of edge cases, and produces inscrutable errors when it goes wrong. I'm happy to be proven wrong (can work out some examples over slack if you're interested) and obviously I'm not a decision-maker on either of these packages, so it's just a friendly tip-off from someone who's been bitten from this kind of thing.

[MasonProtter]

I have a use case for this that I think is better expressed with ReflectOn than the metadata system: I want to take a type Foo{T} and while in a pass, treat it as if Foo{T} <: T.

The problem with doing such a thing via metadata is that if I’m calling f(::Foo{T}), I don’t have an appropriate object of typeT to stick in the first argument.

[vchuravy]

I have a use case for this that I think is better expressed with ReflectOn than the metadata system: I want to take a type Foo{T} and while in a pass, treat it as if Foo{T} <: T.

What do you actually want to achieve?

[oxinabox]

In general I feel like the idea that a small tweak to fix a function that exists but doesn't accept your type is a thing that one wants to do a lot.

It's like a halfway step between normal dispatch, and the tagging system ?

[MasonProtter]

What do you actually want to achieve?

It's for symbolic programming. The idea is to have

abstract type Symbolic{T} end

struct Sym{T}
    name::Symbol
end

struct SymExpr{T}
    op
    args::Vector{Any}
end

and then be able to do things like e.g.

x = Sym{Real}(:x)
f(x::Real) = 1 + x

julia> @symbolic f(x)
SymExpr{Real}((+), [Sym{Real}(:x), 1])

and then I'd like to turn around and be able to do

v = Sym{AbstractVector{<:Complex{Real}}}
g(v::AbstractVector) = v'v

julia> @symbolic g(v)
SymExpr{Complex{Real}}(*, [SymExpr{Adjoint{AbstractVector{Complex{Real}}}}(adjoint, [v]), v])

Or something where I symbolically represent strings, etc. With single inheritance and no pervasive usage of traits in julia, it seems that something like a cassette pass is the only way to achieve this.

For this to work, I'd essentially just need this PR and then I'd have to set up a system to 'register' functions so that doing @register foo will tell the overdub that it's okay to expand into the body of foo, and if it isn't registered then just leave foo as a primitive in the graph.

[MasonProtter]

Hm, after some further reflection (heh), I think the current ReflectOn implementation may be a bit too aggressive for my needs.

Currently, it seems to start a whole new overdub pass, whereas I think a safer interface would be to only descend down a single function call, basically invoke(f, Tuple{Ts...}, args...) except there would be no requirement that xs <: Tuple{Ts...}.

Nope, I was just "holding it wrong", this works great for my purposes.

[shashi]

@MasonProtter thanks for the excellent examples! @YingboMa and I talked about exactly this problem (type of Expressions in ModelingToolkit) but realized we can't do it. It's cool that there is a path to do it!

[shashi]

I think Valentin's concern about general usefulness vs maintenance is warranted.

I think what Valentin was suggesting was to try and remove the Tagging system code. But Valentin, I think that we really need to make it work fast, not remove it. :-p

[MasonProtter]

Has anyone had any thoughts on whether this should or should not be merged? It would be quite useful for enhancing ForwardDiff2 and some of my potential Cassette based projects. It seems to be a more fundamental and lower maintenance facility than something like tagging.