Update Architectures.jl

CliMA · Jan 29, 2025 · bb121c8 · bb121c8
1 parent 0cf9c8f
commit bb121c8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ext/OceananigansReactantExt/Architectures.jl b/ext/OceananigansReactantExt/Architectures.jl
@@ -11,13 +11,13 @@ architecture(::Reactant.AnyTracedRArray) = ReactantState
 array_type(::ReactantState) = ConcreteRArray
 
 on_architecture(::ReactantState, a::Array) = ConcreteRArray(a)
-on_architecture(::ReactantState, a::ConcreteRArray) = a
-on_architecture(::ReactantState, a::TracedRArray) = a
+on_architecture(::ReactantState, a::Reactant.AnyConcreteRArray) = a
+on_architecture(::ReactantState, a::Reactant.AnyTracedRArray) = a
 on_architecture(::ReactantState, a::BitArray) = ConcreteRArray(a)
 on_architecture(::ReactantState, a::SubArray{<:Any, <:Any, <:Array}) = ConcreteRArray(a)
 
 unified_array(::ReactantState, a) = a
 
-@inline device_copy_to!(dst::AnyConcreteRArray, src::AnyConcreteRArray; kw...) = Base.copyto!(dst, src)
+@inline device_copy_to!(dst::Reactant.AnyConcreteRArray, src::Reactant.AnyConcreteRArray; kw...) = Base.copyto!(dst, src)
 
 end # module