Fix unbroadcast_prod_*, again

dfdx · Feb 8, 2020 · 80b4886 · 80b4886
1 parent 4911d5b
commit 80b4886
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/helpers.jl b/src/helpers.jl
@@ -73,11 +73,11 @@ function unbroadcast_prod_x(x::AbstractArray, y::AbstractArray, Δ)
 end
 unbroadcast_prod_y(x::AbstractArray, y::AbstractArray, Δ) = unbroadcast_prod_x(y, x, Δ)
 
-
-unbroadcast_prod_x(x::Number, y::AbstractArray, Δ) = unbroadcast_prod_x(convert(typeof(y), [x]), y, Δ)[1]
-unbroadcast_prod_x(x::AbstractArray, y::Number, Δ) = unbroadcast_prod_x(x, convert(typeof(x), [y]), Δ)
-unbroadcast_prod_y(x::AbstractArray, y::Number, Δ) = unbroadcast_prod_y(x, convert(typeof(x), [y]), Δ)[1]
-unbroadcast_prod_y(x::Number, y::AbstractArray, Δ) = unbroadcast_prod_y(convert(typeof(y), [x]), y, Δ)
+device_like(example, a) = (device = guess_device([example]); device(a))
+unbroadcast_prod_x(x::Number, y::AbstractArray, Δ) = unbroadcast_prod_x(device_like(y, [x]), y, Δ)[1]
+unbroadcast_prod_x(x::AbstractArray, y::Number, Δ) = unbroadcast_prod_x(x, device_like(x, [y]), Δ)
+unbroadcast_prod_y(x::AbstractArray, y::Number, Δ) = unbroadcast_prod_y(x, device_like(x, [y]), Δ)[1]
+unbroadcast_prod_y(x::Number, y::AbstractArray, Δ) = unbroadcast_prod_y(device_like(y, [x]), y, Δ)
 
 
 untranspose_vec(ds::Transpose{T, <:AbstractVector{T}}) where T = transpose(ds)