xlogy broadcast adjoint

2020-05-12 17:29:35 +01:00 · 2020-05-12 17:29:35 +01:00 · f5a8900ffb
parent bd43201f37
commit f5a8900ffb
1 changed files with 5 additions and 0 deletions
--- a/src/layers/stateless.jl
+++ b/src/layers/stateless.jl
@ -288,3 +288,8 @@ CuArrays.@cufunc function xlogy(x, y)
  result = x * log(y)
  ifelse(iszero(x), zero(result), result)
 end
+
+@adjoint function broadcasted(::typeof(xlogy), x::Zygote.Numeric, y::Zygote.Numeric)
+  res = xlogy.(x, y)
+  res, Δ -> (nothing, Zygote.unbroadcast(x, xlogy.(Δ, y)), Zygote.unbroadcast(y, Δ .* x ./ y))
+end