feat: error on common mistakes

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.1.20"
+version = "0.1.21"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

ext/LuxCoreEnzymeCoreExt.jl

@@ -1,9 +1,37 @@
 module LuxCoreEnzymeCoreExt
 
-using EnzymeCore: EnzymeRules
+using EnzymeCore: EnzymeCore, EnzymeRules
 using LuxCore: LuxCore
 using Random: AbstractRNG
 
 EnzymeRules.inactive(::typeof(LuxCore.replicate), ::AbstractRNG) = nothing
 
+# Handle common mistakes users might make
+const LAYER_DERIVATIVE_ERROR_MSG = """
+Lux Layers only support `EnzymeCore.Const` annotation.
+
+Lux Layers are immutable constants and gradients w.r.t. them are `nothing`. To
+compute the gradients w.r.t. the layer's parameters, use the first argument returned
+by `LuxCore.setup(rng, layer)` instead.
+"""
+
+function EnzymeCore.Active(::LuxCore.AbstractExplicitLayer)
+    throw(ArgumentError(LAYER_DERIVATIVE_ERROR_MSG))
+end
+
+for annotation in (:Duplicated, :DuplicatedNoNeed)
+    @eval function EnzymeCore.$(annotation)(
+            ::LuxCore.AbstractExplicitLayer, ::LuxCore.AbstractExplicitLayer)
+        throw(ArgumentError(LAYER_DERIVATIVE_ERROR_MSG))
+    end
+end
+
+for annotation in (:BatchDuplicated, :BatchDuplicatedNoNeed)
+    @eval function EnzymeCore.$(annotation)(
+            ::LuxCore.AbstractExplicitLayer, ::NTuple{N, <:LuxCore.AbstractExplicitLayer},
+            check::Bool=true) where {N}
+        throw(ArgumentError(LAYER_DERIVATIVE_ERROR_MSG))
+    end
+end
+
 end