Add package extensions

mtfishman · mtfishman · commit 09b468ae40a4 · 2025-10-11T11:57:33.000-04:00
diff --git a/NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl b/NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl
@@ -1,7 +1,7 @@
 using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
 using NDTensors.AMDGPUExtensions: AMDGPUExtensions, ROCArrayAdaptor
 using NDTensors.GPUArraysCoreExtensions: storagemode
-using NDTensors.TypeParameterAccessors:
+using NDTensors.Vendored.TypeParameterAccessors:
     default_type_parameters, set_type_parameters, type_parameters
 using Adapt: Adapt, adapt
 using AMDGPU: AMDGPU, ROCArray, ROCVector
diff --git a/NDTensors/ext/NDTensorsAMDGPUExt/linearalgebra.jl b/NDTensors/ext/NDTensorsAMDGPUExt/linearalgebra.jl
@@ -1,7 +1,7 @@
 using NDTensors.AMDGPUExtensions: roc
 using NDTensors.Expose: Expose, Exposed, expose, ql, ql_positive
 using NDTensors.GPUArraysCoreExtensions: cpu
-using NDTensors.TypeParameterAccessors: unwrap_array_type
+using NDTensors.Vendored.TypeParameterAccessors: unwrap_array_type
 using LinearAlgebra: svd
 using Adapt: adapt
 using AMDGPU: ROCMatrix
diff --git a/NDTensors/ext/NDTensorsAMDGPUExt/set_types.jl b/NDTensors/ext/NDTensorsAMDGPUExt/set_types.jl
@@ -1,5 +1,5 @@
 using AMDGPU: ROCArray
 using NDTensors.GPUArraysCoreExtensions: storagemode
-using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
 
 TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(storagemode)) = Position(3)
diff --git a/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl b/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl
@@ -1,4 +1,5 @@
 module NDTensorsCUDAExt
+
 include("append.jl")
 include("default_kwargs.jl")
 include("copyto.jl")
@@ -9,4 +10,5 @@ include("indexing.jl")
 include("linearalgebra.jl")
 include("mul.jl")
 include("permutedims.jl")
+
 end
diff --git a/NDTensors/ext/NDTensorsCUDAExt/adapt.jl b/NDTensors/ext/NDTensorsCUDAExt/adapt.jl
@@ -4,7 +4,7 @@ using Functors: fmap
 using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
 using NDTensors.CUDAExtensions: CUDAExtensions, CuArrayAdaptor
 using NDTensors.GPUArraysCoreExtensions: storagemode
-using NDTensors.TypeParameterAccessors:
+using NDTensors.Vendored.TypeParameterAccessors:
     default_type_parameters, set_type_parameters, type_parameters
 
 function CUDAExtensions.cu(xs; storagemode = default_type_parameters(CuArray, storagemode))
diff --git a/NDTensors/ext/NDTensorsCUDAExt/linearalgebra.jl b/NDTensors/ext/NDTensorsCUDAExt/linearalgebra.jl
@@ -4,7 +4,7 @@ using LinearAlgebra: Adjoint, svd
 using NDTensors: NDTensors
 using NDTensors.Expose: Expose, expose, ql, ql_positive
 using NDTensors.GPUArraysCoreExtensions: cpu
-using NDTensors.TypeParameterAccessors: unwrap_array_type
+using NDTensors.Vendored.TypeParameterAccessors: unwrap_array_type
 function NDTensors.svd_catch_error(A::CuMatrix; alg::String = "jacobi_algorithm")
     if alg == "jacobi_algorithm"
         alg = CUDA.CUSOLVER.JacobiAlgorithm()
diff --git a/NDTensors/ext/NDTensorsCUDAExt/set_types.jl b/NDTensors/ext/NDTensorsCUDAExt/set_types.jl
@@ -1,6 +1,6 @@
 using CUDA: CuArray
 using NDTensors.GPUArraysCoreExtensions: storagemode
-using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
 
 function TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(storagemode))
     return Position(3)
diff --git a/NDTensors/ext/NDTensorsGPUArraysCoreExt/contract.jl b/NDTensors/ext/NDTensorsGPUArraysCoreExt/contract.jl
@@ -2,7 +2,7 @@ using Adapt: adapt
 using GPUArraysCore: AbstractGPUArray
 using NDTensors: NDTensors, DenseTensor, DiagTensor, contract!, dense, inds, Tensor
 using NDTensors.Expose: Exposed, expose, unexpose
-using NDTensors.TypeParameterAccessors: parenttype, set_ndims
+using NDTensors.Vendored.TypeParameterAccessors: parenttype, set_ndims
 
 function NDTensors.contract!(
         output_tensor::Exposed{<:AbstractGPUArray, <:DenseTensor},
diff --git a/NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl b/NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl
@@ -1,7 +1,9 @@
 module NDTensorsJLArraysExt
+
 include("copyto.jl")
 include("indexing.jl")
 include("linearalgebra.jl")
 include("mul.jl")
 include("permutedims.jl")
+
 end
diff --git a/NDTensors/ext/NDTensorsJLArraysExt/linearalgebra.jl b/NDTensors/ext/NDTensorsJLArraysExt/linearalgebra.jl
@@ -4,7 +4,7 @@ using LinearAlgebra: LinearAlgebra, Hermitian, Symmetric, qr, eigen
 using NDTensors: NDTensors
 using NDTensors.Expose: Expose, expose, qr, qr_positive, ql, ql_positive
 using NDTensors.GPUArraysCoreExtensions: cpu
-using NDTensors.TypeParameterAccessors: unwrap_array_type
+using NDTensors.Vendored.TypeParameterAccessors: unwrap_array_type
 
 ## TODO this function exists because of the same issue below. when
 ## that issue is resolved we can rely on the abstractarray version of
diff --git a/NDTensors/ext/NDTensorsMetalExt/adapt.jl b/NDTensors/ext/NDTensorsMetalExt/adapt.jl
@@ -5,7 +5,7 @@ using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
 using NDTensors.Expose: Exposed
 using NDTensors.MetalExtensions: MetalExtensions, MtlArrayAdaptor
 using NDTensors.GPUArraysCoreExtensions: GPUArraysCoreExtensions
-using NDTensors.TypeParameterAccessors: set_type_parameters, type_parameters
+using NDTensors.Vendored.TypeParameterAccessors: set_type_parameters, type_parameters
 
 GPUArraysCoreExtensions.cpu(e::Exposed{<:MtlArray}) = adapt(Array, e)
 
diff --git a/NDTensors/ext/NDTensorsMetalExt/linearalgebra.jl b/NDTensors/ext/NDTensorsMetalExt/linearalgebra.jl
@@ -1,7 +1,7 @@
 using Metal: MtlMatrix
 using LinearAlgebra: LinearAlgebra, qr, eigen, svd
 using NDTensors.Expose: qr_positive, ql_positive, ql
-using NDTensors.TypeParameterAccessors:
+using NDTensors.Vendored.TypeParameterAccessors:
     set_type_parameters, type_parameters, unwrap_array_type
 
 function LinearAlgebra.qr(A::Exposed{<:MtlMatrix})
diff --git a/NDTensors/ext/NDTensorsMetalExt/set_types.jl b/NDTensors/ext/NDTensorsMetalExt/set_types.jl
@@ -1,6 +1,6 @@
 using Metal: MtlArray
 using NDTensors.GPUArraysCoreExtensions: storagemode
-using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
 
 function TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(storagemode))
     return Position(3)
diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl
@@ -1,7 +1,23 @@
 module NDTensors
 
 module Vendored
-    include(joinpath("vendored", "TypeParameterAccessors", "src", "TypeParameterAccessors.jl"))
+    include(
+        joinpath(
+            "vendored", "TypeParameterAccessors", "src", "TypeParameterAccessors.jl"
+        )
+    )
+    include(
+        joinpath(
+            "vendored", "TypeParameterAccessors", "ext",
+            "TypeParameterAccessorsFillArraysExt.jl"
+        )
+    )
+    include(
+        joinpath(
+            "vendored", "TypeParameterAccessors", "ext",
+            "TypeParameterAccessorsStridedViewsExt.jl"
+        )
+    )
 end
 
 #####################################
diff --git a/NDTensors/src/lib/AMDGPUExtensions/src/AMDGPUExtensions.jl b/NDTensors/src/lib/AMDGPUExtensions/src/AMDGPUExtensions.jl
@@ -1,14 +1,5 @@
 module AMDGPUExtensions
 
-module Vendored
-    include(
-        joinpath(
-            "..", "..", "..", "vendored", "TypeParameterAccessors", "src",
-            "TypeParameterAccessors.jl",
-        )
-    )
-end
-
 include("roc.jl")
 
 end
diff --git a/NDTensors/src/lib/AMDGPUExtensions/src/roc.jl b/NDTensors/src/lib/AMDGPUExtensions/src/roc.jl
@@ -1,4 +1,4 @@
-using .Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+using ..Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
 using ..GPUArraysCoreExtensions: storagemode
 # Implemented in NDTensorsAMDGPUExt
 function roc end
diff --git a/NDTensors/src/lib/CUDAExtensions/src/CUDAExtensions.jl b/NDTensors/src/lib/CUDAExtensions/src/CUDAExtensions.jl
@@ -1,14 +1,5 @@
 module CUDAExtensions
 
-module Vendored
-    include(
-        joinpath(
-            "..", "..", "..", "vendored", "TypeParameterAccessors", "src",
-            "TypeParameterAccessors.jl",
-        )
-    )
-end
-
 include("cuda.jl")
 
 end
diff --git a/NDTensors/src/lib/CUDAExtensions/src/cuda.jl b/NDTensors/src/lib/CUDAExtensions/src/cuda.jl
@@ -1,4 +1,4 @@
-using .Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+using ..Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
 using ..GPUArraysCoreExtensions: storagemode
 # Implemented in NDTensorsCUDAExt
 function cu end
diff --git a/NDTensors/src/lib/Expose/src/Expose.jl b/NDTensors/src/lib/Expose/src/Expose.jl
@@ -1,14 +1,5 @@
 module Expose
 
-module Vendored
-    include(
-        joinpath(
-            "..", "..", "..", "vendored", "TypeParameterAccessors", "src",
-            "TypeParameterAccessors.jl",
-        )
-    )
-end
-
 using SimpleTraits
 using LinearAlgebra
 using Base: ReshapedArray
diff --git a/NDTensors/src/lib/Expose/src/exposed.jl b/NDTensors/src/lib/Expose/src/exposed.jl
@@ -1,4 +1,4 @@
-using .Vendored.TypeParameterAccessors:
+using ..Vendored.TypeParameterAccessors:
     TypeParameterAccessors, unwrap_array_type, parenttype, type_parameters
 struct Exposed{Unwrapped, Object}
     object::Object
diff --git a/NDTensors/src/lib/GPUArraysCoreExtensions/src/GPUArraysCoreExtensions.jl b/NDTensors/src/lib/GPUArraysCoreExtensions/src/GPUArraysCoreExtensions.jl
@@ -1,14 +1,5 @@
 module GPUArraysCoreExtensions
 
-module Vendored
-    include(
-        joinpath(
-            "..", "..", "..", "vendored", "TypeParameterAccessors", "src",
-            "TypeParameterAccessors.jl",
-        )
-    )
-end
-
 include("gpuarrayscore.jl")
 
 end
diff --git a/NDTensors/src/lib/GPUArraysCoreExtensions/src/gpuarrayscore.jl b/NDTensors/src/lib/GPUArraysCoreExtensions/src/gpuarrayscore.jl
@@ -1,5 +1,5 @@
 using ..Expose: Exposed, unexpose
-using .Vendored.TypeParameterAccessors: TypeParameterAccessors, type_parameters, set_type_parameters
+using ..Vendored.TypeParameterAccessors: TypeParameterAccessors, type_parameters, set_type_parameters
 
 function storagemode(object)
     return storagemode(typeof(object))
diff --git a/NDTensors/src/lib/MetalExtensions/src/MetalExtensions.jl b/NDTensors/src/lib/MetalExtensions/src/MetalExtensions.jl
@@ -1,14 +1,5 @@
 module MetalExtensions
 
-module Vendored
-    include(
-        joinpath(
-            "..", "..", "..", "vendored", "TypeParameterAccessors", "src",
-            "TypeParameterAccessors.jl",
-        )
-    )
-end
-
 include("metal.jl")
 
 end
diff --git a/NDTensors/src/lib/MetalExtensions/src/metal.jl b/NDTensors/src/lib/MetalExtensions/src/metal.jl
@@ -1,4 +1,4 @@
-using .Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+using ..Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
 using ..GPUArraysCoreExtensions: storagemode
 # Implemented in NDTensorsMetalExt
 function mtl end
diff --git a/NDTensors/src/lib/RankFactorization/src/RankFactorization.jl b/NDTensors/src/lib/RankFactorization/src/RankFactorization.jl
@@ -1,14 +1,5 @@
 module RankFactorization
 
-module Vendored
-    include(
-        joinpath(
-            "..", "..", "..", "vendored", "TypeParameterAccessors", "src",
-            "TypeParameterAccessors.jl",
-        )
-    )
-end
-
 include("default_kwargs.jl")
 include("truncate_spectrum.jl")
 include("spectrum.jl")
diff --git a/NDTensors/src/lib/RankFactorization/src/default_kwargs.jl b/NDTensors/src/lib/RankFactorization/src/default_kwargs.jl
@@ -1,4 +1,4 @@
-using .Vendored.TypeParameterAccessors: unwrap_array_type
+using ..Vendored.TypeParameterAccessors: unwrap_array_type
 replace_nothing(::Nothing, replacement) = replacement
 replace_nothing(value, replacement) = value
 
diff --git a/NDTensors/src/lib/RankFactorization/src/truncate_spectrum.jl b/NDTensors/src/lib/RankFactorization/src/truncate_spectrum.jl
@@ -1,4 +1,4 @@
-using .Vendored.TypeParameterAccessors: unwrap_array_type
+using ..Vendored.TypeParameterAccessors: unwrap_array_type
 
 ## TODO write Exposed version of truncate
 function truncate!!(P::AbstractArray; kwargs...)
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsAMDGPUExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsAMDGPUExt.jl
@@ -0,0 +1,12 @@
+module TypeParameterAccessorsAMDGPUExt
+
+using AMDGPU: AMDGPU, ROCArray
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+
+TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(ndims)) = Position(2)
+function TypeParameterAccessors.default_type_parameters(::Type{<:ROCArray})
+    return (Float64, 1, AMDGPU.Mem.HIPBuffer)
+end
+
+end
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsCUDAExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsCUDAExt.jl
@@ -0,0 +1,12 @@
+module TypeParameterAccessorsCUDAExt
+
+using CUDA: CUDA, CuArray
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+
+TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(ndims)) = Position(2)
+function TypeParameterAccessors.default_type_parameters(::Type{<:CuArray})
+    return (Float64, 1, CUDA.default_memory)
+end
+
+end
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsFillArraysExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsFillArraysExt.jl
@@ -0,0 +1,17 @@
+module TypeParameterAccessorsFillArraysExt
+
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+using FillArrays: Fill, Zeros, Ones
+
+for T in (:Fill, :Zeros, :Ones)
+    @eval begin
+        TypeParameterAccessors.position(::Type{<:$T}, ::typeof(eltype)) = Position(1)
+        TypeParameterAccessors.position(::Type{<:$T}, ::typeof(ndims)) = Position(2)
+        TypeParameterAccessors.position(::Type{<:$T}, ::typeof(axes)) = Position(3)
+        TypeParameterAccessors.default_type_parameters(::Type{<:$T}) = (
+            Float64, 1, Tuple{Base.OneTo{Int}},
+        )
+    end
+end
+
+end
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsJLArraysExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsJLArraysExt.jl
@@ -0,0 +1,10 @@
+module TypeParameterAccessorsJLArraysExt
+
+using JLArrays: JLArray
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+
+TypeParameterAccessors.position(::Type{<:JLArray}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{<:JLArray}, ::typeof(ndims)) = Position(2)
+TypeParameterAccessors.default_type_parameters(::Type{<:JLArray}) = (Float64, 1)
+
+end
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsMetalExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsMetalExt.jl
@@ -0,0 +1,12 @@
+module TypeParameterAccessorsMetalExt
+
+using Metal: Metal, MtlArray
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+
+TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(ndims)) = Position(2)
+function TypeParameterAccessors.default_type_parameters(::Type{<:MtlArray})
+    return (Float64, 1, Metal.DefaultStorageMode)
+end
+
+end
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsStridedViewsExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsStridedViewsExt.jl
@@ -0,0 +1,15 @@
+module TypeParameterAccessorsStridedViewsExt
+
+using StridedViews: StridedView
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position, parenttype
+
+TypeParameterAccessors.position(::Type{<:StridedView}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{<:StridedView}, ::typeof(ndims)) = Position(2)
+function TypeParameterAccessors.position(::Type{<:StridedView}, ::typeof(parenttype))
+    return Position(3)
+end
+function TypeParameterAccessors.default_type_parameters(::Type{<:StridedView})
+    return (Float64, 1, Vector{Float64}, typeof(identity))
+end
+
+end
diff --git a/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsoneAPIExt.jl b/NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsoneAPIExt.jl
@@ -0,0 +1,12 @@
+module TypeParameterAccessorsoneAPIExt
+
+using oneAPI: oneAPI, oneArray
+using NDTensors.Vendored.TypeParameterAccessors: TypeParameterAccessors, Position
+
+TypeParameterAccessors.position(::Type{<:oneArray}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{<:oneArray}, ::typeof(ndims)) = Position(2)
+function TypeParameterAccessors.default_type_parameters(::Type{<:oneAPI})
+    return (Float64, 1, oneAPI.oneL0.DeviceBuffer)
+end
+
+end
diff --git a/NDTensors/vendor/run.jl b/NDTensors/vendor/run.jl
@@ -13,4 +13,6 @@ for (name, version) in deps
     end
     mkpath(joinpath(p, "src"))
     cp(joinpath(local_path, "src"), joinpath(p, "src"); force = true)
+    mkpath(joinpath(p, "ext"))
+    cp(joinpath(local_path, "ext"), joinpath(p, "ext"); force = true)
 end
