Flux: Error when computing gradient of broadcasted ifelse on GPU

Specific Domains Machine Learning
,
1

Here is a code example

X = rand(3,3) .- 0.4
X = X |> cu
f(X) = sum(ifelse.(X .> 0, X, 1f0))
gradient(f, X)

And here is the error message

ERROR: TypeError: non-boolean (ForwardDiff.Dual{Nothing,Bool,3}) used in boolean context

If X is a CPU array, then it works fine, but not when X is a CuArray.
Any suggestion on fixing the issue?

Thanks

2

That ifelse looks somewhat suspect, try extracting it as an element-wise function:

weirdrelu(x) = ifelse(x > 0, x, one(x))
f(X) = sum(weirdrelu.(X))
gradient(f, X)

Returns ([1, 1, 1],) as expected on the latest CUDA.jl.