[NDTensors] JLArrays Extension

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.34"
+version = "0.3.35"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -36,6 +36,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
@@ -47,6 +48,7 @@ NDTensorsAMDGPUExt = ["AMDGPU", "GPUArraysCore"]
 NDTensorsCUDAExt = ["CUDA", "GPUArraysCore"]
 NDTensorsGPUArraysCoreExt = "GPUArraysCore"
 NDTensorsHDF5Ext = "HDF5"
+NDTensorsJLArraysExt = ["GPUArraysCore", "JLArrays"]
 NDTensorsMappedArraysExt = ["MappedArrays"]
 NDTensorsMetalExt = ["GPUArraysCore", "Metal"]
 NDTensorsOctavianExt = "Octavian"
@@ -70,6 +72,7 @@ GPUArraysCore = "0.1"
 HDF5 = "0.14, 0.15, 0.16, 0.17"
 HalfIntegers = "1"
 InlineStrings = "1"
+JLArrays = "0.1"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"
 MappedArrays = "0.4"
@@ -95,6 +98,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"

NDTensors/ext/NDTensorsAMDGPUExt/set_types.jl

@@ -6,6 +6,5 @@ using AMDGPU: AMDGPU, ROCArray
 function TypeParameterAccessors.default_type_parameters(::Type{<:ROCArray})
   return (Float64, 1, AMDGPU.Mem.HIPBuffer)
 end
-TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(eltype)) = Position(1)
-TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(ndims)) = Position(2)
+
 TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(storagemode)) = Position(3)

NDTensors/ext/NDTensorsCUDAExt/set_types.jl

@@ -3,12 +3,6 @@ using CUDA: CUDA, CuArray
 using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
 using NDTensors.GPUArraysCoreExtensions: storagemode
 
-function TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(eltype))
-  return Position(1)
-end
-function TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(ndims))
-  return Position(2)
-end
 function TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(storagemode))
   return Position(3)
 end

NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl

@@ -0,0 +1,3 @@
+module NDTensorsJLArraysExt
+include("set_types.jl")
+end

NDTensors/ext/NDTensorsJLArraysExt/set_types.jl

@@ -0,0 +1,7 @@
+# TypeParameterAccessors definitions
+using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
+using JLArrays: JLArray
+
+function TypeParameterAccessors.default_type_parameters(::Type{<:JLArray})
+  return (Float64, 1)
+end

NDTensors/ext/NDTensorsMetalExt/set_types.jl

@@ -4,13 +4,6 @@ using Metal: Metal, MtlArray
 using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
 using NDTensors.GPUArraysCoreExtensions: storagemode
 
-## TODO remove TypeParameterAccessors when SetParameters is removed
-function TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(eltype))
-  return Position(1)
-end
-function TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(ndims))
-  return Position(2)
-end
 function TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(storagemode))
   return Position(3)
 end

NDTensors/test/NDTensorsTestUtils/device_list.jl

@@ -1,5 +1,7 @@
-using NDTensors: NDTensors
 using Pkg: Pkg
+using JLArrays
+using NDTensors: NDTensors
+
 if "cuda" in ARGS || "all" in ARGS
   ## Right now adding CUDA during Pkg.test results in a
   ## compat issues. I am adding it back to test/Project.toml
@@ -28,6 +30,7 @@ function devices_list(test_args)
   devs = Vector{Function}(undef, 0)
   if isempty(test_args) || "base" in test_args
     push!(devs, NDTensors.cpu)
+    # push!(devs, jl)
   end
 
   if "cuda" in test_args || "cutensor" in test_args || "all" in test_args
@@ -47,5 +50,6 @@ function devices_list(test_args)
   if "metal" in test_args || "all" in test_args
     push!(devs, NDTensors.MetalExtensions.mtl)
   end
+
   return devs
 end

NDTensors/test/Project.toml

@@ -1,14 +1,14 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
@@ -29,5 +29,6 @@ cuTENSOR = "2.0"
 
 [extras]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"