NNlib Support

[willtebbutt]

#169 highlights that we need rules for a range of functionality which lives in NNlib.jl -- this is not surprising, and has largely already been done (see e.g. the Enzyme extension). Someone needs to systematically work through and test that Tapir works on everything in NNlib.

[yebai]

Can we re-use the ChainRules implementation? See, e.g. the rule for dropout.

[willtebbutt]

In short: yes. We don't have a nice macro for doing it at the minute, but it ought to be quite doable provided that we constrain the types that impose some of our own constraints on the types.

In the case of dropout: this one actually might be a little tricky because it uses kwargs, and I've not looked at supporting kwargs yet (this largely happen automatically), but it should otherwise be fine.

Something like batched_transpose should be quite straightforward. You would just do something along the lines of

function rrule!!(::CoDual{typeof(batched_transpose)}, A::CoDual{Array{<:Any,3}})
    B, pb = rrule(batched_transpose, primal(A)
    function pb!!(dB)
        _, dA_inc = pb(dB)
        increment!!(tangent(A), dA_inc)
        return NoRData(), NoRData()
    end
    return zero_fcodual(B), pb!!
end

Note that I've restricted the array type here -- we would just need to be careful with how we extend this rule to other array types, as it's only valid in Tapir.jl for arrays whose tangents are arrays, rather than Tangents.

[yebai]

Note that I've restricted the array type here -- we would just need to be careful with how we extend this rule to other array types, as it's only valid in Tapir.jl for arrays whose tangents are arrays, rather than Tangents.

Thanks @willtebbutt. A follow-up question on this topic: if we put mutation aside, is it true in general, ChainRules rules are a superset of Tapir rules, in the sense that one can always get a Tapir rule by restricting argument types on a ChainRules rule.

[willtebbutt]

In the narrow sense that lots of ChainRules.rrules are written with abstract arguments: yes.

However, recall that this gives you an overly optimistic sense of how much more general a ChainRules.rrule is. It's not very hard to construct e.g. subtypes of AbstractArray{<:Any, 3} for which this won't work.

[yebai]

In the narrow sense that lots of ChainRules.rrules are written with abstract arguments: yes.

Was the current ChainRules design motivated by something, or is it an oversight?

[willtebbutt]

A bit of both.

It's motivated by trying to write rules for things like Zygote, where you really want your rule system to do a lot of work, because Zygote can't differentiate much itself due to lack of mutation support.

It's an oversight in the sense that we didn't realise how hard it is to write rules which work for lots of types that are robust. The upshot is the kinds of issues we've discussed previously (mainly around poor composition performance).

[yebai]

But ChainRules doesn't prevent one from writing rules with concrete argument types in principle. Is that correct?

Or does writing rules with concrete argument types always require something like typed tangents? IIUC, the typed tangents in Tapirs seem like an extension of ChainRules via concrete types rather than something incompatible.

[willtebbutt]

But ChainRules doesn't prevent one from writing rules with concrete argument types in principle. Is that correct?

Correct. Equally, Tapir.jl doesn't prevent you from writing rules with abstract arguments types, it's just not generally a good idea.

Or does writing rules with concrete argument types always require something like typed tangents? IIUC, the typed tangents in Tapirs seem like an extension of ChainRules via concrete types rather than something incompatible.

I'm not sure what you mean by "typed tangents" -- could you expand?

[yebai]

I'm not sure what you mean by "typed tangents" -- could you expand?

I am referring to the fact that, in Tapir, each primal type has a unique tangent type. Do writing ChainRule rules with concrete argument types always require this unique tangent-type assumption?

[willtebbutt]

Hmm okay. Could you provide a small code example showing what you mean? I'm still not quite following what you're trying to say / ask, so I think I need some examples.

[willtebbutt]

This is partially addressed by #254, but there remains work to be done on GPU integration, and generalising the functionality added to AbstractArrays other than Arrays. Consequently, I'm going to keep this issue open.