feat: emit batch norm training op

LuxDL · Feb 11, 2025 · 5f1abca · 5f1abca
1 parent 044efef
commit 5f1abca
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.6.1"
+version = "1.7.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

diff --git a/lib/LuxLib/ext/LuxLibReactantExt.jl b/lib/LuxLib/ext/LuxLibReactantExt.jl
@@ -2,7 +2,7 @@ module LuxLibReactantExt
 
 using Reactant: Reactant, MLIR, Ops, TracedUtils, TracedRArray, AnyTracedRArray,
                 AnyTracedRVector, TracedRNumber
-using Static: False
+using Static: True, False
 
 using LuxLib: LuxLib, Impl, Optional, Utils
 
@@ -56,14 +56,11 @@ function Impl.batchnorm(
     return act.(TracedRArray{T, ndims(x)}((), res, size(x))), rμ, rσ²
 end
 
-# The following code is commented out since we don't have Batchnorm Op Adjoint registered
-# for EnzymeJAX yet
-#=
 function Impl.batchnorm(
         x::AnyTracedRArray{T},
         γ::Optional{<:AnyTracedRVector}, β::Optional{<:AnyTracedRVector},
         rμ::Optional{<:AnyTracedRVector}, rσ²::Optional{<:AnyTracedRVector},
-        training::StaticBool, act::F, momentum, ϵ
+        ::True, act::F, momentum, ϵ
 ) where {T, F}
     x = TracedUtils.materialize_traced_array(x)
 
@@ -101,6 +98,5 @@ function Impl.batchnorm(
         return res, rμ, rσ²
     end
 end
-=#
 
 end