use in softmax too

Author: mcabbott

src/softmax.jl

@@ -69,7 +69,7 @@ function softmax!(out::AbstractArray{T}, x::AbstractArray; dims = 1) where {T}
 end
 
 function ∇softmax_data(dy::AbstractArray{T}, y::AbstractArray{S}; dims = 1) where {T,S}
-    dx = if within_grad()
+    dx = if within_gradient(y)
         tmp = dy .* y
         tmp .- y .* sum(tmp; dims)
     else
@@ -88,9 +88,6 @@ function rrule(::typeof(softmax), x; dims = 1)
     return y, softmax_pullback
 end
 
-within_grad() = false
-rrule(::typeof(within_grad)) = true, _ -> (NoTangent(),)
-
 fast_maximum(x::AbstractArray{T}; dims) where {T} = @fastmath reduce(max, x; dims, init = float(T)(-Inf))
 
 """