refactor: cleanup some old pre-1.0 hacks (#1102)

LuxDL · Nov 24, 2024 · cd96335 · cd96335 · github-actions · Nov 24, 2024
1 parent d755929
commit cd96335
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 13 deletions.
diff --git a/src/extended_ops.jl b/src/extended_ops.jl
@@ -224,19 +224,18 @@ CRC.@non_differentiable istraining(::Any)
 
 end
 
-using .LuxOps: LuxOps, multigate
+using .LuxOps: LuxOps, multigate, xlogx, xlogy, foldl_init
 
 const safe_getproperty = LuxOps.getproperty
 const safe_eachslice = LuxOps.eachslice
 
-# TODO: directly import them from LuxOps from 1.0
-const private_xlogx = LuxOps.xlogx
-const private_xlogy = LuxOps.xlogy
-const private_foldl_init = LuxOps.foldl_init
-
 # These are defined here to avoid a circular dependency among modules
-for (op, field) in (:bias => :use_bias, :affine => :affine,
-    :track_stats => :track_stats, :train_state => :train_state)
+for (op, field) in (
+    :bias => :use_bias,
+    :affine => :affine,
+    :track_stats => :track_stats,
+    :train_state => :train_state
+)
     @eval function $(Symbol(:has_, op))(l::AbstractLuxLayer)
         res = known(safe_getproperty(l, Val($(Meta.quot(field)))))
         return ifelse(res === nothing, false, res)

diff --git a/src/helpers/losses.jl b/src/helpers/losses.jl
@@ -261,8 +261,7 @@ end
 
 for logits in (true, false)
     return_expr = logits ? :(return loss.agg((1 .- ỹ) .* ŷ .- logsigmoid.(ŷ))) :
-                  :(return loss.agg(-private_xlogy.(ỹ, ŷ .+ ϵ) .-
-                                    private_xlogy.(1 .- ỹ, 1 .- ŷ .+ ϵ)))
+                  :(return loss.agg(-xlogy.(ỹ, ŷ .+ ϵ) .- xlogy.(1 .- ỹ, 1 .- ŷ .+ ϵ)))
 
     @eval function unsafe_apply_loss(loss::BinaryCrossEntropyLoss{$(logits)}, ŷ, y)
         T = promote_type(eltype(ŷ), eltype(y))
@@ -387,7 +386,7 @@ for logits in (true, false)
                   :(return LossFunctionImpl.fused_agg(
         loss.agg, -, sum(ỹ .* logsoftmax(ŷ; loss.dims); loss.dims))) :
                   :(return LossFunctionImpl.fused_agg(
-        loss.agg, -, sum(private_xlogy.(ỹ, ŷ .+ ϵ); loss.dims)))
+        loss.agg, -, sum(xlogy.(ỹ, ŷ .+ ϵ); loss.dims)))
 
     @eval function unsafe_apply_loss(loss::CrossEntropyLoss{$(logits)}, ŷ, y)
         T = promote_type(eltype(ŷ), eltype(y))
@@ -603,7 +602,7 @@ end
 function unsafe_apply_loss(loss::KLDivergenceLoss, ŷ, y)
     cross_entropy = unsafe_apply_loss(loss.celoss, ŷ, y)
     # Intentional broadcasting for Zygote type stability
-    entropy = loss.agg(sum(private_xlogx.(y); loss.dims))
+    entropy = loss.agg(sum(xlogx.(y); loss.dims))
     return entropy + cross_entropy
 end
 

diff --git a/src/layers/recurrent.jl b/src/layers/recurrent.jl
@@ -131,7 +131,7 @@ function (r::Recurrence{True})(x::Union{AbstractVector, NTuple}, ps, st::NamedTu
         (out, carry), state = apply(r.cell, (input, carry), ps, state)
         return vcat(outputs, [out]), carry, state
     end
-    results = private_foldl_init(recur_op, x)
+    results = foldl_init(recur_op, x)
     return first(results), last(results)
 end