Where to find basic aritmetic operators derivatives and broadcasted versions

[Cvikli]

Hey,

I try to redefine @scalar_rules and some of the macros, to create the appropriate code for our symbolic derivation.

I tried to google and understand the source code it but didin't see where to fing the pullback for

:(+)(::AbstractArray, ::AbstractArray) ...
:(.+)(::AbstractArray, ::AbstractArray) ...
... -,*,/,^...

and so on.

I know it is easy to do but the library looks really nice and I don't understand where are the basic aritmetics. I found some *... in some case but the Array*Array also something I couldn't find.

What do I miss? How does Zygote do the chainrules without these arithmetics I couldn't get to know. Can you guys help me?

[nickrobinson251]

The frule and rrules for scalar-functions are mostly defined with the @scalar_rule macro.

And rules for scalar + and * are in the src/rulesets/base/fastmath_able.jl file.

See

ChainRules.jl/src/rulesets/Base/fastmath_able.jl

Lines 163 to 165 in 38caf4b

	@scalar_rule x + y (One(), One())
	@scalar_rule x - y (One(), -1)
	@scalar_rule x / y (one(x) / y, -(Ω / y))

and

ChainRules.jl/src/rulesets/Base/fastmath_able.jl

Lines 206 to 220 in 38caf4b

	# product rule requires special care for arguments where `mul` is non-commutative
	function frule((_, Δx, Δy), ::typeof(*), x::Number, y::Number)
	# Optimized version of `Δx .* y .+ x .* Δy`. Also, it is potentially more
	# accurate on machines with FMA instructions, since there are only two
	# rounding operations, one in `muladd/fma` and the other in `*`.
	∂xy = muladd.(Δx, y, x .* Δy)
	return x * y, ∂xy
	end
	function rrule(::typeof(*), x::Number, y::Number)
	function times_pullback(ΔΩ)
	return (NO_FIELDS, ΔΩ * y', x' * ΔΩ)
	end
	return x * y, times_pullback
	end

Rules for Array functions are in rulesets/base/array.jl or rulesets/base/arraymath.jl (roughly trying to match the location of the functions in Julia Base). Some Array functionality is from the LinearAlgebra standard library, so defined in src/rulesets/LinearAlgebra/.

The rules for Array*Array are here

ChainRules.jl/src/rulesets/Base/arraymath.jl

Lines 18 to 63 in 76ef95c

	#####
	##### `*`
	#####
	function rrule(
	::typeof(*),
	A::AbstractVecOrMat{<:CommutativeMulNumber},
	B::AbstractVecOrMat{<:CommutativeMulNumber},
	)
	function times_pullback(Ȳ)
	return (
	NO_FIELDS,
	InplaceableThunk(
	@thunk(Ȳ * B'),
	X̄ -> mul!(X̄, Ȳ, B', true, true)
	),
	InplaceableThunk(
	@thunk(A' * Ȳ),
	X̄ -> mul!(X̄, A', Ȳ, true, true)
	)
	)
	end
	return A * B, times_pullback
	end
	function rrule(
	::typeof(*),
	A::AbstractVector{<:CommutativeMulNumber},
	B::AbstractMatrix{<:CommutativeMulNumber},
	)
	function times_pullback(Ȳ)
	@assert size(B, 1) === 1 # otherwise primal would have failed.
	return (
	NO_FIELDS,
	InplaceableThunk(
	@thunk(Ȳ * vec(B')),
	X̄ -> mul!(X̄, Ȳ, vec(B'), true, true)
	),
	InplaceableThunk(
	@thunk(A' * Ȳ),
	X̄ -> mul!(X̄, A', Ȳ, true, true)
	)
	)
	end
	return A * B, times_pullback
	end

[Cvikli]

Thank you for the detailed answer!

For me the unclear part is that how does * handle for bigger arrays 3D, 4D... But as I typed this again realised that there is no * opreation between bigger array, what I am looking for the ".*" and ".^" etc. which is interpreted between bigger arrays. What are the broadcasted function's version frule, rrule?

One more question just asking fast, I see * but only see @scalar_rule for + and -, from where does the Zygote get the info for the array cases?

[simeonschaub]

Zygote doesn't use ChainRules for handling broadcasting at all, everything is defined here: https://github.com/FluxML/Zygote.jl/blob/master/src/lib/broadcast.jl. We'd first need to solve JuliaDiff/ChainRulesCore.jl#68 before we can define such rules in ChainRules.

[Cvikli]

It is really tricky code, I just can't understand how does this broadcasting called/overloaded in each . case. Also I don't see the cases for the N dimension Array types. But I will try to read more of their code maybe.

So the only question I have where are the array versions of the +, - in ChainRules, or am I missing something here?

[ToucheSir]

So the only question I have where are the array versions of the +, - in ChainRules, or am I missing something here?

See

ChainRules.jl/src/rulesets/Base/arraymath.jl

Lines 306 to 329 in 2e6491c

	#####
	##### Negation (Unary -)
	#####
	function rrule(::typeof(-), x::AbstractArray)
	function negation_pullback(ȳ)
	return NoTangent(), InplaceableThunk(ā -> ā .-= ȳ, @thunk(-ȳ))
	end
	return -x, negation_pullback
	end
	#####
	##### Addition (Multiarg `+`)
	#####
	function rrule(::typeof(+), arrs::AbstractArray...)
	y = +(arrs...)
	arr_axs = map(axes, arrs)
	function add_pullback(dy)
	return (NoTangent(), map(ax -> reshape(dy, ax), arr_axs)...)
	end
	return y, add_pullback
	end

. I'm not sure if there's a dedicated rrule for -(A::AbstractArray, B::AbstractArray) or if it falls through to broadcasting.