Unable to differentiate a broadcasted constructor with CUDA array inputs

[BioTurboNick]

Due to a CUDA bug: JuliaGPU/CUDA.jl#2514, crossposted for reference.

Encountered when differentiating over:

using Distributions
using CUDA
a = cu(ones(5)); b = cu(zeros(5));

Normal.(a, b)

[BioTurboNick]

It seems the CUDA behavior is not likely to change, so it would need to be addressed in Zygote.

The issue occurs because broadcast_forward wraps a function around the broadcasted constructor to dualize it.

[ToucheSir]

IIRC this is a fundamental problem with the Type/constructor duality (UnionAlls) in Julia, and would be quite tricky for Zygote to address on its own. The Turing folks ran into a similar issue with broadcasting Distribution constructors, and they solved it by using an intermediate function which doesn't capture the constructor type. e.g.:

make_normal(a, b) = Normal(a, b)
# a = cu(...); b = cu(...)
make_normal.(a, b)

In terms of things to try, one could look at whether specializing

Zygote.jl/src/lib/broadcast.jl

Lines 273 to 278 in 2517e67

	@inline function dual_function(f::F) where F
	function (args::Vararg{Any,N}) where N
	ds = dualize(args...)
	return f(ds...)
	end
	end

on Type removes the capture and thus the (non-zero size) closure.