Issue with forward AD

[ebelnikola]

Hi!

I tried to do forward differentiation with tensors from TensorKit and I noticed that I cannot properly take the derivative of something as simple as t*A with respect to t at zero. Here, A is a tensor. After a short investigation I noticed that the problem lies in this package and in particular in the behavior of the scale function applied to numbers. It is currently implemented as follows:

@inline scale(x::Number, α::Number) = (iszero(α) ? zero(x) : x) * α

The problem appears when one tries to call scale(x, α) with α=Dual{Float64}(0.0,1.0) and gets Dual{Float64}(0.0,0.0) instead of Dual{Float64}(0.0,x). This happens because iszero(α)==true, which is reasonable. Thus, for purposes of forward AD it would be nice to either remove this check or to add some frule. As for the latter solution, I suspect it will not work for Zygote and ForwardDiff as it seems that Zygote uses ForwardDiff for forward differentiation and the latter relies on DiffRules instead of ChainRules.

[lkdvos]

Hi!

Thanks for reporting this (apologies it took me some time to get to it, this somehow slipped through the cracks).
As far as I remember, there is actually a good reason for this implementation, which has to do with how NaN values are handled.

julia> x = NaN;

julia> α = 0.0;

julia> scale(x, α)
0.0

julia> x * α
NaN

However, I actually think that this is one of the reasons we introduced VectorInterface.Zero, to actually make that a type-domain check, so perhaps it could be reasonable to simply remove the if statement. I am a bit scared though about hidden bugs in downstream code that might depend on this behavior and is not explicitly using Zero() for that.

Do you know which rule system we would have to overload to make your case work? Is it just a matter of defining this for dual numbers, or is there a more general way of tackling this?