xlogy

src/layers/stateless.jl

@@ -54,15 +54,15 @@ function huber_loss(ŷ, y;  δ=eltype(ŷ)(1))
 end
 
 function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Nothing)
-  return -sum(y .* log.(ŷ)) * 1 // size(y, 2)
+  return -sum(xlogy.(y, ŷ)) * 1 // size(y, 2)
 end
 
 function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Number)
-  return -sum(y .* log.(ŷ)) .* weight * 1 // size(y, 2)
+  return -sum(xlogy.(y, ŷ)) .* weight * 1 // size(y, 2)
 end
 
 function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::AbstractVector)
-  return -sum(y .* log.(ŷ) .* weight) * 1 // size(y, 2)
+  return -sum(xlogy.(y, ŷ) .* weight) * 1 // size(y, 2)
 end
 
 """
@@ -123,7 +123,7 @@ julia> Flux.binarycrossentropy.(σ.([-1.1491, 0.8619, 0.3127]), [1, 1, 0])
  0.8616703662235441
 ```
 """
-binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)
+binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -xlogy(y, ŷ + ϵ) - xlogy(1 - y, 1 - ŷ + ϵ)
 
 # Re-definition to fix interaction with CuArrays.
 CuArrays.@cufunc binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)
@@ -195,7 +195,7 @@ It is always non-negative and zero only when both the distributions are equal
 everywhere.
 """
 function kldivergence(ŷ, y)
-  entropy = sum(y .* log.(y)) * 1 //size(y,2)
+  entropy = sum(xlogx.(y)) * 1 //size(y,2)
   cross_entropy = crossentropy(ŷ, y)
   return entropy + cross_entropy
 end
@@ -208,7 +208,7 @@ distribution `y`; calculated as `sum(ŷ .- y .* log.(ŷ)) / size(y, 2)`.
 
 [More information.](https://peltarion.com/knowledge-center/documentation/modeling-view/build-an-ai-model/loss-functions/poisson).
 """
-poisson(ŷ, y) = sum(ŷ .- y .* log.(ŷ)) * 1 // size(y,2)
+poisson(ŷ, y) = sum(ŷ .- xlogy.(y, ŷ)) * 1 // size(y,2)
 
 """
     hinge(ŷ, y)
@@ -262,3 +262,16 @@ by linearizing all values for each element in the batch.
 function flatten(x::AbstractArray)
   return reshape(x, :, size(x)[end])
 end
+
+"""
+    xlogx(x::Real)
+Return `x * log(x)` for `x ≥ 0`, handling `x = 0` by taking the downward limit.
+"""
+xlogx(x::Real) = x > zero(x) ? x * log(x) : zero(log(x))
+
+"""
+    xlogy(x::Real, y::Real)
+Return `x * log(y)` for `y > 0` with correct limit at `x = 0`.
+"""
+xlogy(x::T, y::T) where {T<:Real} = x > zero(T) ? x * log(y) : zero(log(x))
+xlogy(x::Real, y::Real) = xlogy(promote(x, y)...)

test/layers/stateless.jl

@@ -1,9 +1,18 @@
 using Test
 using Flux: onehotbatch, mse, crossentropy, logitcrossentropy,
-            σ, binarycrossentropy, logitbinarycrossentropy, flatten
+            σ, binarycrossentropy, logitbinarycrossentropy, flatten,
+            xlogx, xlogy
 
 const ϵ = 1e-7
 
+@testset "xlogx & xlogy" begin
+    @test iszero(xlogx(0))
+    @test xlogx(2) ≈ 2.0 * log(2.0)
+
+    @test iszero(xlogy(0, 1))
+    @test xlogy(2, 3) ≈ 2.0 * log(3.0)
+end
+
 @testset "losses" begin
   # First, regression-style y's
   y = [1, 1, 0, 0]