Merge pull request #302 from aussetg/adapt_wrappers

Adapt rules for batched_* wrappers
FluxML · Mar 23, 2021 · a9727e4 · a9727e4 · CarloLucibello · Mar 23, 2021
2 parents aa5a6bd + 3e417ed
commit a9727e4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/Project.toml b/Project.toml
@@ -3,6 +3,7 @@ uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 version = "0.7.18"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

diff --git a/src/batched/batchedadjtrans.jl b/src/batched/batchedadjtrans.jl
@@ -1,6 +1,7 @@
 using LinearAlgebra
 
 import Base: -
+import Adapt: adapt_structure, adapt
 
 _batched_doc = """
     batched_transpose(A::AbstractArray{T,3})
@@ -100,3 +101,6 @@ function rrule(::typeof(batched_adjoint), A::AbstractArray{<:Any,3})
     b_adjoint_back(Δ) = (NO_FIELDS, batched_adjoint(Δ))
     batched_adjoint(A), b_adjoint_back
 end
+
+adapt_structure(to, x::BatchedAdjoint) = BatchedAdjoint(adapt(to, parent(x)))
+adapt_structure(to, x::BatchedTranspose) = BatchedTranspose(adapt(to, parent(x)))