Type assert for safety

LuxDL · Jun 23, 2024 · 66a5cda · 66a5cda
1 parent 2b00266
commit 66a5cda
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/ext/LuxReverseDiffExt/training.jl b/ext/LuxReverseDiffExt/training.jl
@@ -18,15 +18,15 @@ end
 # Uncompiled ReverseDiff
 @inline function __uncompiled_reverse_diff(obj_fn::F, data, ts::TrainState) where {F}
  grads = Lux.recursive_make_zero(ts.parameters)
- ts_new = TrainState(TrainingBackendCache{:ReverseDiff, true}(grads, nothing),
- obj_fn, ts.model, ts.parameters, ts.states, ts.optimizer_state, ts.step)
+ ts_new = TrainState(TrainingBackendCache{:ReverseDiff, true}(grads, nothing), obj_fn,
+ ts.model, ts.parameters, ts.states, ts.optimizer_state, ts.step)
  return __uncompiled_reverse_diff(obj_fn, data, ts_new)
 end
 
 @inline function __uncompiled_reverse_diff(obj_fn::F, data,
  ts::TrainState{<:TrainingBackendCache{:ReverseDiff, FT}}) where {F, FT}
- tape = ReverseDiff.InstructionTape()
  dparams = FT ? ts.cache.dparameters : Lux.recursive_make_zero!!(ts.cache.dparameters)
+ tape = ReverseDiff.InstructionTape()
  ps_tracked = Lux.recursive_map(
  Lux.__Fix3(ReverseDiff.TrackedArray, tape), ts.parameters, dparams)
 
@@ -44,12 +44,22 @@ end
 # Compiled ReverseDiff
 @inline function __compiled_reverse_diff(obj_fn::F, data, ts::TrainState) where {F}
  grads = Lux.recursive_make_zero(ts.parameters)
- ts_new = TrainState(TrainingBackendCache{:ReverseDiff, true}(grads, nothing),
- obj_fn, ts.model, ts.parameters, ts.states, ts.optimizer_state, ts.step)
+ ts_new = TrainState(TrainingBackendCache{:ReverseDiff, true}(grads, nothing), obj_fn,
+ ts.model, ts.parameters, ts.states, ts.optimizer_state, ts.step)
  return __compiled_reverse_diff(obj_fn, data, ts_new)
 end
 
-@inline function __compiled_reverse_diff(obj_fn::F, data,
- ts::TrainState{<:TrainingBackendCache{:ReverseDiff, FT}}) where {F, FT}
+## Tape hasn't been compiled yet
+@inline function __compiled_reverse_diff(obj_fn::F,
+ data,
+ ts::TrainState{<:TrainingBackendCache{:ReverseDiff, FT, P, Nothing}}) where {
+ F, FT, P}
+ dparams = FT ? ts.cache.dparameters : Lux.recursive_make_zero!!(ts.cache.dparameters)
+ tape = ReverseDiff.InstructionTape()
+ ps_tracked = Lux.recursive_map(
+ Lux.__Fix3(ReverseDiff.TrackedArray, tape), ts.parameters, dparams)
+
+ loss, st, stats = obj_fn(ts.model, ps_tracked, ts.states, data)
+
  error(1)
 end
diff --git a/src/helpers/recursive_ops.jl b/src/helpers/recursive_ops.jl
@@ -7,7 +7,7 @@ common cases.
 
 Any leaves of `x` that are arrays and allow in-place addition will be modified in place.
 """
-recursive_add!!(x, y) = recursive_map(__add!!, x, y)
+@inline recursive_add!!(x, y) = recursive_map(__add!!, x, y)
 
 """
  recursive_eltype(x)
@@ -38,7 +38,7 @@ Recursively create a zero value for a nested structure `x`. This is equivalent t
 
 See also [`Lux.recursive_make_zero!!`](@ref).
 """
-@inline recursive_make_zero(x) = recursive_map(__zero, x)
+@inline recursive_make_zero(x) = recursive_map(__zero, x)::typeof(x)
 
 """
  recursive_make_zero!!(x)
@@ -48,7 +48,7 @@ in-place zeroing will be modified in place.
 
 See also [`Lux.recursive_make_zero`](@ref) for fully out-of-place version.
 """
-@inline recursive_make_zero!!(x) = recursive_map(__zero!!, x)
+@inline recursive_make_zero!!(x) = recursive_map(__zero!!, x)::typeof(x)
 
 """
  recursive_map(f, x, args...)