No commit message

paulxshen · Jul 20, 2024 · d15ea2f · d15ea2f
1 parent 703c868
commit d15ea2f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,14 +1,16 @@
 name = "ArrayPadding"
 uuid = "fe5ef43c-4912-423d-9170-4bd5cb6a013a"
 authors = ["Paul Shen <pxshen@alumni.stanford.edu>"]
-version = "0.1.25"
+version = "0.1.26"
 
 [deps]
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Lazy = "50d2b5c4-7a5e-59d5-8109-a42b560f39c0"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 Functors = "0.4"
 Lazy = "0.15"
 julia = "1.9"
+GPUArraysCore="0.1,0.2"
diff --git a/src/alg.jl b/src/alg.jl
@@ -1,4 +1,4 @@
-function lr(a, b, i, l, r, out=true, ol=0, or=ol)
+function lr(a::T, b, i, l, r, out=true, ol=0, or=ol) where {T}
     d = ndims(a)
     al = ar = nothing
     if l > 0
@@ -68,7 +68,12 @@ function lr(a, b, i, l, r, out=true, ol=0, or=ol)
             ar = fill(b, [j == i ? r : size(a, j) for j = 1:d]...)
         end
     end
-    al, ar
+    if T <: AbstractGPUArray
+        return map((al, ar)) do x
+            isnothing(x) ? nothing : T(x)
+        end
+    end
+    return (al, ar)
 end
 
 Base.vec(x::Number, d) = fill(x, d)

diff --git a/src/main.jl b/src/main.jl
@@ -1,4 +1,4 @@
-using Functors, Statistics
+using Functors, Statistics, GPUArraysCore
 using Lazy: @forward
 # using Zygote: Buffer, bufferfrom
 include("types.jl")