Add iscu function

ITensor · Oct 8, 2023 · 46e00ce · 46e00ce
1 parent 0f310e9
commit 46e00ce
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl b/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl
@@ -17,6 +17,8 @@ else
 end
 
 include("imports.jl")
+
+iscu(::Type{<:CuArray}) = true
 include("set_types.jl")
 include("adapt.jl")
 include("linearalgebra.jl")

diff --git a/NDTensors/ext/NDTensorsCUDAExt/imports.jl b/NDTensors/ext/NDTensorsCUDAExt/imports.jl
@@ -1,6 +1,6 @@
 import NDTensors: cu, set_ndims, set_eltype, set_eltype_if_unspecified, similartype
 import NDTensors:
   ContractionProperties, _contract!, GemmBackend, auto_select_backend, _gemm!
-import NDTensors.SetParameters: nparameters, get_parameter, set_parameter, default_parameter
+import NDTensors.SetParameters: nparameters, get_parameter, set_parameter, default_parameter, iscu
 
 import .CUDA: CuArrayAdaptor
diff --git a/NDTensors/src/abstractarray/similar.jl b/NDTensors/src/abstractarray/similar.jl
@@ -56,6 +56,9 @@ function similar(arraytype::Type{<:AbstractArray}, dims::Tuple)
   return similartype(arraytype, shape)(undef, NDTensors.to_shape(arraytype, shape))
 end
 
+# For when there are CUArray specific issues inline
+iscu(A::AbstractArray) = iscu(leaf_parenttype(A))
+iscu(::Type{<:AbstractArray}) = false
 # This function actually allocates the data.
 # Catches conversions of dimensions specified by ranges
 # dimensions specified by integers with `Base.to_shape`.