From 3b670826664bb46ee0e8cd7711b8c99bc1fb9c92 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 10 Oct 2025 21:08:38 -0400 Subject: [PATCH 1/7] [WIP] [NDTensors] Vendor TypeParameterAccessors.jl --- NDTensors/Project.toml | 2 +- NDTensors/vendor/Project.toml | 5 +++++ NDTensors/vendor/README.md | 13 +++++++++++++ NDTensors/vendor/run.jl | 16 ++++++++++++++++ 4 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 NDTensors/vendor/Project.toml create mode 100644 NDTensors/vendor/README.md create mode 100644 NDTensors/vendor/run.jl diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 7df9c48cdc..085be6bf96 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -1,7 +1,7 @@ name = "NDTensors" uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" authors = ["Matthew Fishman "] -version = "0.4.13" +version = "0.4.14" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" diff --git a/NDTensors/vendor/Project.toml b/NDTensors/vendor/Project.toml new file mode 100644 index 0000000000..4c7dfde138 --- /dev/null +++ b/NDTensors/vendor/Project.toml @@ -0,0 +1,5 @@ +[deps] +PackageAnalyzer = "e713c705-17e4-4cec-abe0-95bf5bf3e10c" + +[compat] +PackageAnalyzer = "3.1" diff --git a/NDTensors/vendor/README.md b/NDTensors/vendor/README.md new file mode 100644 index 0000000000..ebd43458f0 --- /dev/null +++ b/NDTensors/vendor/README.md @@ -0,0 +1,13 @@ +# vendor + +Here we vendor some of our dependencies to load our own copies of them. This means: + +- the functions defined in them are not shared with the "real" package they come from, so we do not see method extensions of them +- we have our own private copy of the module, allowing us to load a different version than the user + +To update the dependencies, install the `vendor` environment and invoke `run.jl`: + +```sh +julia --project=vendor -e 'using Pkg; Pkg.instantiate()' +julia --project=vendor vendor/run.jl +``` diff --git a/NDTensors/vendor/run.jl b/NDTensors/vendor/run.jl new file mode 100644 index 0000000000..2491b21f5a --- /dev/null +++ b/NDTensors/vendor/run.jl @@ -0,0 +1,16 @@ +using PackageAnalyzer: PackageAnalyzer + +deps = [("TypeParameterAccessors", v"0.3.11")] + +for (name, version) in deps + pkg = PackageAnalyzer.find_package(name; version) + local_path, reachable, _ = PackageAnalyzer.obtain_code(pkg) + @assert reachable + p = mkpath(joinpath(@__DIR__, "..", "src", "vendored", name)) + # remove any existing files + if isdir(p) + rm(p; recursive = true, force = true) + end + mkpath(joinpath(p, "src")) + cp(joinpath(local_path, "src"), joinpath(p, "src"); force = true) +end From 40dc0a348cd569a4ef06410f1295f3597efcc235 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 10 Oct 2025 21:12:23 -0400 Subject: [PATCH 2/7] Update usings --- NDTensors/Project.toml | 2 -- NDTensors/src/abstractarray/generic_array_constructors.jl | 2 +- NDTensors/src/abstractarray/iscu.jl | 2 +- NDTensors/src/abstractarray/set_types.jl | 2 +- NDTensors/src/abstractarray/similar.jl | 2 +- NDTensors/src/adapt.jl | 2 +- NDTensors/src/blocksparse/blockdims.jl | 2 +- NDTensors/src/blocksparse/blocksparsetensor.jl | 6 +++--- NDTensors/src/blocksparse/diagblocksparse.jl | 2 +- NDTensors/src/blocksparse/linearalgebra.jl | 2 +- NDTensors/src/blocksparse/similar.jl | 2 +- NDTensors/src/default_kwargs.jl | 2 +- NDTensors/src/dense/dense.jl | 2 +- NDTensors/src/dense/generic_array_constructors.jl | 2 +- NDTensors/src/dense/set_types.jl | 2 +- NDTensors/src/diag/diagtensor.jl | 2 +- NDTensors/src/diag/set_types.jl | 2 +- NDTensors/src/diag/similar.jl | 2 +- NDTensors/src/dims.jl | 2 +- NDTensors/src/empty/empty.jl | 4 ++-- NDTensors/src/lib/AMDGPUExtensions/src/roc.jl | 2 +- NDTensors/src/lib/CUDAExtensions/src/cuda.jl | 2 +- NDTensors/src/lib/Expose/src/exposed.jl | 2 +- .../src/lib/GPUArraysCoreExtensions/src/gpuarrayscore.jl | 2 +- NDTensors/src/lib/MetalExtensions/src/metal.jl | 2 +- NDTensors/src/lib/RankFactorization/src/default_kwargs.jl | 2 +- .../src/lib/RankFactorization/src/truncate_spectrum.jl | 2 +- NDTensors/src/linearalgebra/linearalgebra.jl | 2 +- NDTensors/src/linearalgebra/svd.jl | 2 +- NDTensors/src/tensor/set_types.jl | 2 +- NDTensors/src/tensor/similar.jl | 2 +- NDTensors/src/tensorstorage/default_storage.jl | 2 +- NDTensors/src/tensorstorage/set_types.jl | 2 +- NDTensors/src/tensorstorage/similar.jl | 2 +- NDTensors/src/truncate.jl | 4 ++-- 35 files changed, 38 insertions(+), 40 deletions(-) diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 085be6bf96..e5edea4365 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -28,7 +28,6 @@ Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67" StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143" TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" -TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138" VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8" [weakdeps] @@ -88,7 +87,6 @@ StridedViews = "0.2.2, 0.3, 0.4" TBLIS = "0.2" TimerOutputs = "0.5.5" TupleTools = "1.2.0" -TypeParameterAccessors = "0.3" VectorInterface = "0.4.2, 0.5" cuTENSOR = "2" julia = "1.10" diff --git a/NDTensors/src/abstractarray/generic_array_constructors.jl b/NDTensors/src/abstractarray/generic_array_constructors.jl index 6c82102e48..2bf99eba92 100644 --- a/NDTensors/src/abstractarray/generic_array_constructors.jl +++ b/NDTensors/src/abstractarray/generic_array_constructors.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: +using .Vendored.TypeParameterAccessors: unwrap_array_type, specify_default_type_parameters, specify_type_parameters, diff --git a/NDTensors/src/abstractarray/iscu.jl b/NDTensors/src/abstractarray/iscu.jl index e6fc96d8c9..185f7295a9 100644 --- a/NDTensors/src/abstractarray/iscu.jl +++ b/NDTensors/src/abstractarray/iscu.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type # TODO: Make `isgpu`, `ismtl`, etc. # For `isgpu`, will require a `NDTensorsGPUArrayCoreExt`. iscu(A::AbstractArray) = iscu(typeof(A)) diff --git a/NDTensors/src/abstractarray/set_types.jl b/NDTensors/src/abstractarray/set_types.jl index c17f732c34..2e393392f6 100644 --- a/NDTensors/src/abstractarray/set_types.jl +++ b/NDTensors/src/abstractarray/set_types.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors """ # Do we still want to define things like this? diff --git a/NDTensors/src/abstractarray/similar.jl b/NDTensors/src/abstractarray/similar.jl index 23b8329204..9e0cea080e 100644 --- a/NDTensors/src/abstractarray/similar.jl +++ b/NDTensors/src/abstractarray/similar.jl @@ -1,5 +1,5 @@ using Base: DimOrInd, Dims, OneTo -using TypeParameterAccessors: IsWrappedArray, unwrap_array_type, set_eltype, similartype +using .Vendored.TypeParameterAccessors: IsWrappedArray, unwrap_array_type, set_eltype, similartype ## Custom `NDTensors.similar` implementation. ## More extensive than `Base.similar`. diff --git a/NDTensors/src/adapt.jl b/NDTensors/src/adapt.jl index 2244a5c2d1..ba449bcee4 100644 --- a/NDTensors/src/adapt.jl +++ b/NDTensors/src/adapt.jl @@ -27,7 +27,7 @@ double_precision(x) = fmap(x -> adapt(double_precision(eltype(x)), x), x) # Used to adapt `EmptyStorage` types # -using TypeParameterAccessors: specify_type_parameters +using .Vendored.TypeParameterAccessors: specify_type_parameters function adapt_storagetype(to::Type{<:AbstractVector}, x::Type{<:TensorStorage}) return set_datatype(x, specify_type_parameters(to, eltype, eltype(x))) end diff --git a/NDTensors/src/blocksparse/blockdims.jl b/NDTensors/src/blocksparse/blockdims.jl index a73af992a1..f4cfe89de4 100644 --- a/NDTensors/src/blocksparse/blockdims.jl +++ b/NDTensors/src/blocksparse/blockdims.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors """ BlockDim diff --git a/NDTensors/src/blocksparse/blocksparsetensor.jl b/NDTensors/src/blocksparse/blocksparsetensor.jl index 1bb9fe69fc..89a8011d68 100644 --- a/NDTensors/src/blocksparse/blocksparsetensor.jl +++ b/NDTensors/src/blocksparse/blocksparsetensor.jl @@ -1,5 +1,5 @@ using SparseArrays: nnz -using TypeParameterAccessors: similartype +using .Vendored.TypeParameterAccessors: similartype # # BlockSparseTensor (Tensor using BlockSparse storage) @@ -258,7 +258,7 @@ end # Defaults to adding zeros. # Returns the offset of the new block added. # XXX rename to insertblock!, no need to return offset -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type using .Expose: Exposed, expose, unexpose function insertblock_offset!(T::BlockSparseTensor{ElT, N}, newblock::Block{N}) where {ElT, N} newdim = blockdim(T, newblock) @@ -792,7 +792,7 @@ end # permfactor(perm, block, inds) = 1 -using TypeParameterAccessors: set_type_parameters, parenttype +using .Vendored.TypeParameterAccessors: set_type_parameters, parenttype function permutedims!( R::BlockSparseTensor{<:Number, N}, T::BlockSparseTensor{<:Number, N}, diff --git a/NDTensors/src/blocksparse/diagblocksparse.jl b/NDTensors/src/blocksparse/diagblocksparse.jl index 4fe47ec603..2b52c6434b 100644 --- a/NDTensors/src/blocksparse/diagblocksparse.jl +++ b/NDTensors/src/blocksparse/diagblocksparse.jl @@ -1,5 +1,5 @@ using LinearAlgebra: LinearAlgebra -using TypeParameterAccessors: similartype +using .Vendored.TypeParameterAccessors: similartype export DiagBlockSparse, DiagBlockSparseTensor diff --git a/NDTensors/src/blocksparse/linearalgebra.jl b/NDTensors/src/blocksparse/linearalgebra.jl index 8310f14a87..6bcafb8c9b 100644 --- a/NDTensors/src/blocksparse/linearalgebra.jl +++ b/NDTensors/src/blocksparse/linearalgebra.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type using .Expose: expose const BlockSparseMatrix{ElT, StoreT, IndsT} = BlockSparseTensor{ElT, 2, StoreT, IndsT} const DiagBlockSparseMatrix{ElT, StoreT, IndsT} = DiagBlockSparseTensor{ElT, 2, StoreT, IndsT} diff --git a/NDTensors/src/blocksparse/similar.jl b/NDTensors/src/blocksparse/similar.jl index 46ba2ecae2..a850710e0d 100644 --- a/NDTensors/src/blocksparse/similar.jl +++ b/NDTensors/src/blocksparse/similar.jl @@ -1,5 +1,5 @@ using SparseArrays: nnz -using TypeParameterAccessors: similartype +using .Vendored.TypeParameterAccessors: similartype # NDTensors.similar function similar(storagetype::Type{<:BlockSparse}, blockoffsets::BlockOffsets, dims::Tuple) diff --git a/NDTensors/src/default_kwargs.jl b/NDTensors/src/default_kwargs.jl index fda4fe9786..9cb170b0e5 100644 --- a/NDTensors/src/default_kwargs.jl +++ b/NDTensors/src/default_kwargs.jl @@ -1,4 +1,4 @@ -using 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/dense/dense.jl b/NDTensors/src/dense/dense.jl index 8c2a5fec31..abf18e4a47 100644 --- a/NDTensors/src/dense/dense.jl +++ b/NDTensors/src/dense/dense.jl @@ -124,7 +124,7 @@ dense(storagetype::Type{<:Dense}) = storagetype # TODO: make these more general, move to tensorstorage.jl datatype(storetype::Type{<:Dense{<:Any, DataT}}) where {DataT} = DataT -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type function promote_rule( ::Type{<:Dense{ElT1, DataT1}}, ::Type{<:Dense{ElT2, DataT2}} ) where {ElT1, DataT1, ElT2, DataT2} diff --git a/NDTensors/src/dense/generic_array_constructors.jl b/NDTensors/src/dense/generic_array_constructors.jl index 98b2375f0e..a34a97dcba 100644 --- a/NDTensors/src/dense/generic_array_constructors.jl +++ b/NDTensors/src/dense/generic_array_constructors.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: +using .Vendored.TypeParameterAccessors: default_type_parameters, parenttype, set_eltype, diff --git a/NDTensors/src/dense/set_types.jl b/NDTensors/src/dense/set_types.jl index 21d6c9c4e0..8ac590f003 100644 --- a/NDTensors/src/dense/set_types.jl +++ b/NDTensors/src/dense/set_types.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors, Position, parenttype +using .Vendored.TypeParameterAccessors: TypeParameterAccessors, Position, parenttype function set_datatype(storagetype::Type{<:Dense}, datatype::Type{<:AbstractVector}) return Dense{eltype(datatype), datatype} diff --git a/NDTensors/src/diag/diagtensor.jl b/NDTensors/src/diag/diagtensor.jl index f10f78eee7..de9db5b3e9 100644 --- a/NDTensors/src/diag/diagtensor.jl +++ b/NDTensors/src/diag/diagtensor.jl @@ -113,7 +113,7 @@ function dense(::Type{<:Tensor{ElT, N, StoreT, IndsT}}) where {ElT, N, StoreT <: return Tensor{ElT, N, dense(StoreT), IndsT} end -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type # convert to Dense function dense(T::DiagTensor) R = zeros(dense(typeof(T)), inds(T)) diff --git a/NDTensors/src/diag/set_types.jl b/NDTensors/src/diag/set_types.jl index b908710e40..4527136471 100644 --- a/NDTensors/src/diag/set_types.jl +++ b/NDTensors/src/diag/set_types.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors function TypeParameterAccessors.set_eltype(storagetype::Type{<:UniformDiag}, eltype::Type) return Diag{eltype, eltype} diff --git a/NDTensors/src/diag/similar.jl b/NDTensors/src/diag/similar.jl index 398e02b9c0..9dcec8dd3c 100644 --- a/NDTensors/src/diag/similar.jl +++ b/NDTensors/src/diag/similar.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors # NDTensors.similar function similar(storagetype::Type{<:Diag}, dims::Dims) diff --git a/NDTensors/src/dims.jl b/NDTensors/src/dims.jl index 81d3cef8f5..14948841a2 100644 --- a/NDTensors/src/dims.jl +++ b/NDTensors/src/dims.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors export dense, dims, dim, mindim, diaglength diff --git a/NDTensors/src/empty/empty.jl b/NDTensors/src/empty/empty.jl index 78aaea0e91..37ca1dcda2 100644 --- a/NDTensors/src/empty/empty.jl +++ b/NDTensors/src/empty/empty.jl @@ -1,5 +1,5 @@ using SparseArrays: SparseArrays -using TypeParameterAccessors: TypeParameterAccessors, set_eltype, similartype +using .Vendored.TypeParameterAccessors: TypeParameterAccessors, set_eltype, similartype # # Represents a tensor order that could be set to any order. @@ -95,6 +95,6 @@ function Base.show(io::IO, mime::MIME"text/plain", S::EmptyStorage) return println(io, typeof(S)) end -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors TypeParameterAccessors.parenttype(empty::Type{<:EmptyStorage}) = storagetype(empty) zero(empty::EmptyStorage) = empty diff --git a/NDTensors/src/lib/AMDGPUExtensions/src/roc.jl b/NDTensors/src/lib/AMDGPUExtensions/src/roc.jl index a3ba228b51..0aa1a58693 100644 --- a/NDTensors/src/lib/AMDGPUExtensions/src/roc.jl +++ b/NDTensors/src/lib/AMDGPUExtensions/src/roc.jl @@ -1,4 +1,4 @@ -using 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/cuda.jl b/NDTensors/src/lib/CUDAExtensions/src/cuda.jl index acaabcea57..60dad8d5d3 100644 --- a/NDTensors/src/lib/CUDAExtensions/src/cuda.jl +++ b/NDTensors/src/lib/CUDAExtensions/src/cuda.jl @@ -1,4 +1,4 @@ -using 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/exposed.jl b/NDTensors/src/lib/Expose/src/exposed.jl index 64030fc64d..3ebc88e416 100644 --- a/NDTensors/src/lib/Expose/src/exposed.jl +++ b/NDTensors/src/lib/Expose/src/exposed.jl @@ -1,4 +1,4 @@ -using 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/gpuarrayscore.jl b/NDTensors/src/lib/GPUArraysCoreExtensions/src/gpuarrayscore.jl index e98d6c70fb..1e37904042 100644 --- 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 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/metal.jl b/NDTensors/src/lib/MetalExtensions/src/metal.jl index 109c0d0f26..3253c7ed66 100644 --- a/NDTensors/src/lib/MetalExtensions/src/metal.jl +++ b/NDTensors/src/lib/MetalExtensions/src/metal.jl @@ -1,4 +1,4 @@ -using 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/default_kwargs.jl b/NDTensors/src/lib/RankFactorization/src/default_kwargs.jl index fda4fe9786..9cb170b0e5 100644 --- a/NDTensors/src/lib/RankFactorization/src/default_kwargs.jl +++ b/NDTensors/src/lib/RankFactorization/src/default_kwargs.jl @@ -1,4 +1,4 @@ -using 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 index d980623364..d05c2b2ebc 100644 --- a/NDTensors/src/lib/RankFactorization/src/truncate_spectrum.jl +++ b/NDTensors/src/lib/RankFactorization/src/truncate_spectrum.jl @@ -1,4 +1,4 @@ -using 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/linearalgebra/linearalgebra.jl b/NDTensors/src/linearalgebra/linearalgebra.jl index f24f02c22e..bc128940e0 100644 --- a/NDTensors/src/linearalgebra/linearalgebra.jl +++ b/NDTensors/src/linearalgebra/linearalgebra.jl @@ -387,7 +387,7 @@ function qr_positive(M::AbstractMatrix) return (Q, R) end -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type """ ql_positive(M::AbstractMatrix) diff --git a/NDTensors/src/linearalgebra/svd.jl b/NDTensors/src/linearalgebra/svd.jl index bf7729e9ff..1dee094836 100644 --- a/NDTensors/src/linearalgebra/svd.jl +++ b/NDTensors/src/linearalgebra/svd.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type # The state of the `svd_recursive` algorithm. function svd_recursive_state(S::AbstractArray, thresh::Float64) return svd_recursive_state(unwrap_array_type(S), S, thresh) diff --git a/NDTensors/src/tensor/set_types.jl b/NDTensors/src/tensor/set_types.jl index 438921b3ed..66c3688f39 100644 --- a/NDTensors/src/tensor/set_types.jl +++ b/NDTensors/src/tensor/set_types.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors, Position, parenttype +using .Vendored.TypeParameterAccessors: TypeParameterAccessors, Position, parenttype function TypeParameterAccessors.set_ndims(arraytype::Type{<:Tensor}, ndims) # TODO: Implement something like: # ```julia diff --git a/NDTensors/src/tensor/similar.jl b/NDTensors/src/tensor/similar.jl index 4ef609b34e..1a64348cad 100644 --- a/NDTensors/src/tensor/similar.jl +++ b/NDTensors/src/tensor/similar.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors, set_indstype, similartype +using .Vendored.TypeParameterAccessors: TypeParameterAccessors, set_indstype, similartype # NDTensors.similar similar(tensor::Tensor) = setstorage(tensor, similar(storage(tensor))) diff --git a/NDTensors/src/tensorstorage/default_storage.jl b/NDTensors/src/tensorstorage/default_storage.jl index 1467244e15..c67a210f55 100644 --- a/NDTensors/src/tensorstorage/default_storage.jl +++ b/NDTensors/src/tensorstorage/default_storage.jl @@ -3,7 +3,7 @@ default_datatype(eltype::Type = default_eltype()) = Vector{eltype} default_eltype() = Float64 -using TypeParameterAccessors: specify_default_type_parameters +using .Vendored.TypeParameterAccessors: specify_default_type_parameters ## TODO use multiple dispace to make this pick between dense and blocksparse function default_storagetype(datatype::Type{<:AbstractArray}, inds::Tuple) datatype = specify_default_type_parameters(datatype) diff --git a/NDTensors/src/tensorstorage/set_types.jl b/NDTensors/src/tensorstorage/set_types.jl index cfba32f65d..87583be20b 100644 --- a/NDTensors/src/tensorstorage/set_types.jl +++ b/NDTensors/src/tensorstorage/set_types.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors +using .Vendored.TypeParameterAccessors: TypeParameterAccessors function TypeParameterAccessors.set_ndims(arraytype::Type{<:TensorStorage}, ndims::Int) # TODO: Change to this once `TensorStorage` types support wrapping # non-AbstractVector types. diff --git a/NDTensors/src/tensorstorage/similar.jl b/NDTensors/src/tensorstorage/similar.jl index 7f1aa5f01f..5c12a3e321 100644 --- a/NDTensors/src/tensorstorage/similar.jl +++ b/NDTensors/src/tensorstorage/similar.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: TypeParameterAccessors, set_ndims, similartype +using .Vendored.TypeParameterAccessors: TypeParameterAccessors, set_ndims, similartype # NDTensors.similar similar(storage::TensorStorage) = setdata(storage, NDTensors.similar(data(storage))) diff --git a/NDTensors/src/truncate.jl b/NDTensors/src/truncate.jl index 5f46bb2655..73fefafdcf 100644 --- a/NDTensors/src/truncate.jl +++ b/NDTensors/src/truncate.jl @@ -1,4 +1,4 @@ -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type ## TODO write Exposed version of truncate function truncate!!(P::AbstractArray; kwargs...) return truncate!!(unwrap_array_type(P), P; kwargs...) @@ -10,7 +10,7 @@ function truncate!!(::Type{<:Array}, P::AbstractArray; kwargs...) return P, truncerr, docut end -using TypeParameterAccessors: unwrap_array_type +using .Vendored.TypeParameterAccessors: unwrap_array_type # GPU fallback version, convert to CPU. function truncate!!(::Type{<:AbstractArray}, P::AbstractArray; kwargs...) P_cpu = cpu(P) From d0b64d1598a29a3b9ddc6118077dfb298df3ac20 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 10 Oct 2025 21:22:41 -0400 Subject: [PATCH 3/7] Vendor TypeParameterAccessors --- NDTensors/src/NDTensors.jl | 5 + .../AMDGPUExtensions/src/AMDGPUExtensions.jl | 10 + .../lib/CUDAExtensions/src/CUDAExtensions.jl | 10 + NDTensors/src/lib/Expose/src/Expose.jl | 10 + .../src/GPUArraysCoreExtensions.jl | 10 + .../MetalExtensions/src/MetalExtensions.jl | 10 + .../src/RankFactorization.jl | 11 + .../src/TypeParameterAccessors.jl | 23 ++ .../src/base/abstractarray.jl | 94 +++++ .../src/base/linearalgebra.jl | 25 ++ .../src/base/similartype.jl | 93 +++++ .../TypeParameterAccessors/src/ndims.jl | 6 + .../src/type_parameters.jl | 331 ++++++++++++++++++ .../TypeParameterAccessors/src/type_utils.jl | 14 + NDTensors/test/Project.toml | 1 - 15 files changed, 652 insertions(+), 1 deletion(-) create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/TypeParameterAccessors.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/base/abstractarray.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/base/linearalgebra.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/base/similartype.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/ndims.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/type_parameters.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/src/type_utils.jl diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index d984c09758..fbb1a320bd 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -1,4 +1,9 @@ module NDTensors + +module Vendored + include(joinpath("vendored", "TypeParameterAccessors", "src", "TypeParameterAccessors.jl")) +end + ##################################### # Imports and exports # diff --git a/NDTensors/src/lib/AMDGPUExtensions/src/AMDGPUExtensions.jl b/NDTensors/src/lib/AMDGPUExtensions/src/AMDGPUExtensions.jl index e9e77e6cc5..2007c8f1ef 100644 --- a/NDTensors/src/lib/AMDGPUExtensions/src/AMDGPUExtensions.jl +++ b/NDTensors/src/lib/AMDGPUExtensions/src/AMDGPUExtensions.jl @@ -1,4 +1,14 @@ module AMDGPUExtensions + +module Vendored + include( + joinpath( + "..", "..", "..", "vendored", "TypeParameterAccessors", "src", + "TypeParameterAccessors.jl", + ) + ) +end + include("roc.jl") end diff --git a/NDTensors/src/lib/CUDAExtensions/src/CUDAExtensions.jl b/NDTensors/src/lib/CUDAExtensions/src/CUDAExtensions.jl index 63a030f9cb..e9ae01be1b 100644 --- a/NDTensors/src/lib/CUDAExtensions/src/CUDAExtensions.jl +++ b/NDTensors/src/lib/CUDAExtensions/src/CUDAExtensions.jl @@ -1,4 +1,14 @@ module CUDAExtensions + +module Vendored + include( + joinpath( + "..", "..", "..", "vendored", "TypeParameterAccessors", "src", + "TypeParameterAccessors.jl", + ) + ) +end + include("cuda.jl") end diff --git a/NDTensors/src/lib/Expose/src/Expose.jl b/NDTensors/src/lib/Expose/src/Expose.jl index fb64927d7b..297fffd4ae 100644 --- a/NDTensors/src/lib/Expose/src/Expose.jl +++ b/NDTensors/src/lib/Expose/src/Expose.jl @@ -1,4 +1,14 @@ 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/GPUArraysCoreExtensions/src/GPUArraysCoreExtensions.jl b/NDTensors/src/lib/GPUArraysCoreExtensions/src/GPUArraysCoreExtensions.jl index 6845816fca..9f6972403b 100644 --- a/NDTensors/src/lib/GPUArraysCoreExtensions/src/GPUArraysCoreExtensions.jl +++ b/NDTensors/src/lib/GPUArraysCoreExtensions/src/GPUArraysCoreExtensions.jl @@ -1,4 +1,14 @@ module GPUArraysCoreExtensions + +module Vendored + include( + joinpath( + "..", "..", "..", "vendored", "TypeParameterAccessors", "src", + "TypeParameterAccessors.jl", + ) + ) +end + include("gpuarrayscore.jl") end diff --git a/NDTensors/src/lib/MetalExtensions/src/MetalExtensions.jl b/NDTensors/src/lib/MetalExtensions/src/MetalExtensions.jl index 504f13f7aa..deadddfbdf 100644 --- a/NDTensors/src/lib/MetalExtensions/src/MetalExtensions.jl +++ b/NDTensors/src/lib/MetalExtensions/src/MetalExtensions.jl @@ -1,4 +1,14 @@ module MetalExtensions + +module Vendored + include( + joinpath( + "..", "..", "..", "vendored", "TypeParameterAccessors", "src", + "TypeParameterAccessors.jl", + ) + ) +end + include("metal.jl") end diff --git a/NDTensors/src/lib/RankFactorization/src/RankFactorization.jl b/NDTensors/src/lib/RankFactorization/src/RankFactorization.jl index a2b950ac93..d41bbbc317 100644 --- a/NDTensors/src/lib/RankFactorization/src/RankFactorization.jl +++ b/NDTensors/src/lib/RankFactorization/src/RankFactorization.jl @@ -1,5 +1,16 @@ module RankFactorization + +module Vendored + include( + joinpath( + "..", "..", "..", "vendored", "TypeParameterAccessors", "src", + "TypeParameterAccessors.jl", + ) + ) +end + include("default_kwargs.jl") include("truncate_spectrum.jl") include("spectrum.jl") + end diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/TypeParameterAccessors.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/TypeParameterAccessors.jl new file mode 100644 index 0000000000..0a51102f6c --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/TypeParameterAccessors.jl @@ -0,0 +1,23 @@ +module TypeParameterAccessors + +# Exports +export type_parameters, get_type_parameters +export nparameters, is_parameter_specified +export default_type_parameters +export set_type_parameters, set_default_type_parameters +export specify_type_parameters, specify_default_type_parameters +export unspecify_type_parameters + +# Imports +using SimpleTraits: SimpleTraits, @traitdef, @traitimpl + +include("type_utils.jl") +include("type_parameters.jl") + +# Implementations +include("ndims.jl") +include("base/abstractarray.jl") +include("base/similartype.jl") +include("base/linearalgebra.jl") + +end diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/base/abstractarray.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/base/abstractarray.jl new file mode 100644 index 0000000000..f45131e688 --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/base/abstractarray.jl @@ -0,0 +1,94 @@ +struct Self end +position(a, ::Self) = Position(0) +position(::Type, ::Self) = Position(0) +function set_type_parameters(type::Type, ::Self, param) + return error("Can't set the parent type of an unwrapped array type.") +end + +position(::Type{AbstractArray}, ::typeof(eltype)) = Position(1) +position(::Type{AbstractArray}, ::typeof(ndims)) = Position(2) +default_type_parameters(::Type{AbstractArray}) = (Float64, 1) + +position(::Type{<:Array}, ::typeof(eltype)) = Position(1) +position(::Type{<:Array}, ::typeof(ndims)) = Position(2) +default_type_parameters(::Type{<:Array}) = (Float64, 1) + +position(::Type{<:BitArray}, ::typeof(ndims)) = Position(1) +default_type_parameters(::Type{<:BitArray}) = (1,) + +function set_eltype(array::AbstractArray, param) + return convert(set_eltype(typeof(array), param), array) +end + +## This will fail if position of `ndims` is not defined for `type` +function set_ndims(type::Type{<:AbstractArray}, param) + return set_type_parameters(type, ndims, param) +end +function set_ndims(type::Type{<:AbstractArray}, param::NDims) + return set_type_parameters(type, ndims, ndims(param)) +end + +# Trait indicating if the AbstractArray type is an array wrapper. +# Assumes that it implements `NDTensors.parenttype`. +@traitdef IsWrappedArray{ArrayType} + +#! format: off +@traitimpl IsWrappedArray{ArrayType} <- is_wrapped_array(ArrayType) +#! format: on + +parenttype(type::Type{<:AbstractArray}) = type_parameters(type, parenttype) +parenttype(object::AbstractArray) = parenttype(typeof(object)) +position(::Type{<:AbstractArray}, ::typeof(parenttype)) = Self() + +is_wrapped_array(arraytype::Type{<:AbstractArray}) = (parenttype(arraytype) ≠ arraytype) +@inline is_wrapped_array(array::AbstractArray) = is_wrapped_array(typeof(array)) +@inline is_wrapped_array(object) = false + +using SimpleTraits: Not, @traitfn + +@traitfn function unwrap_array_type( + arraytype::Type{ArrayType} + ) where {{ArrayType; IsWrappedArray{ArrayType}}} + return unwrap_array_type(parenttype(arraytype)) +end + +@traitfn function unwrap_array_type( + arraytype::Type{ArrayType} + ) where {{ArrayType; !IsWrappedArray{ArrayType}}} + return arraytype +end + +# For working with instances. +unwrap_array_type(array::AbstractArray) = unwrap_array_type(typeof(array)) + +function set_parenttype(t::Type, param) + return set_type_parameters(t, parenttype, param) +end + +@traitfn function set_eltype( + type::Type{ArrayType}, param + ) where {{ArrayType <: AbstractArray; IsWrappedArray{ArrayType}}} + new_parenttype = set_eltype(parenttype(type), param) + # Need to set both in one `set_type_parameters` call to avoid + # conflicts in type parameter constraints of certain wrapper types. + return set_type_parameters(type, (eltype, parenttype), (param, new_parenttype)) +end + +@traitfn function set_eltype( + type::Type{ArrayType}, param + ) where {{ArrayType <: AbstractArray; !IsWrappedArray{ArrayType}}} + return set_type_parameters(type, eltype, param) +end + +for wrapper in [:PermutedDimsArray, :(Base.ReshapedArray), :SubArray] + @eval begin + position(type::Type{<:$wrapper}, ::typeof(eltype)) = Position(1) + position(type::Type{<:$wrapper}, ::typeof(ndims)) = Position(2) + end +end +for wrapper in [:(Base.ReshapedArray), :SubArray] + @eval position(type::Type{<:$wrapper}, ::typeof(parenttype)) = Position(3) +end +for wrapper in [:PermutedDimsArray] + @eval position(type::Type{<:$wrapper}, ::typeof(parenttype)) = Position(5) +end diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/base/linearalgebra.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/base/linearalgebra.jl new file mode 100644 index 0000000000..a7dfde2d35 --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/base/linearalgebra.jl @@ -0,0 +1,25 @@ +using LinearAlgebra: + Adjoint, + Diagonal, + Hermitian, + LowerTriangular, + Symmetric, + Transpose, + UnitLowerTriangular, + UnitUpperTriangular, + UpperTriangular + +for wrapper in [ + :Transpose, + :Adjoint, + :Symmetric, + :Hermitian, + :UpperTriangular, + :LowerTriangular, + :UnitUpperTriangular, + :UnitLowerTriangular, + :Diagonal, + ] + @eval position(::Type{<:$wrapper}, ::typeof(eltype)) = Position(1) + @eval position(::Type{<:$wrapper}, ::typeof(parenttype)) = Position(2) +end diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/base/similartype.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/base/similartype.jl new file mode 100644 index 0000000000..082c269bbb --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/base/similartype.jl @@ -0,0 +1,93 @@ +""" +`set_indstype` should be overloaded for +types with structured dimensions, +like `OffsetArrays` or named indices +(such as ITensors). +""" +function set_indstype(arraytype::Type{<:AbstractArray}, dims::Tuple) + return set_ndims(arraytype, NDims(length(dims))) +end + +function similartype(arraytype::Type{<:AbstractArray}, eltype::Type, ndims::NDims) + return similartype(similartype(arraytype, eltype), ndims) +end + +function similartype(arraytype::Type{<:AbstractArray}, eltype::Type, dims::Tuple) + return similartype(similartype(arraytype, eltype), dims) +end + +@traitfn function similartype( + arraytype::Type{ArrayT} + ) where {{ArrayT; !IsWrappedArray{ArrayT}}} + return arraytype +end + +@traitfn function similartype( + arraytype::Type{ArrayT}, eltype::Type + ) where {{ArrayT; !IsWrappedArray{ArrayT}}} + return set_eltype(arraytype, eltype) +end + +function similartype(arraytype::Type{<:BitArray}, elt::Type) + if is_parameter_specified(arraytype, ndims) + return Array{elt, ndims(arraytype)} + end + return Array{elt} +end + +@traitfn function similartype( + arraytype::Type{ArrayT}, dims::Tuple + ) where {{ArrayT; !IsWrappedArray{ArrayT}}} + return set_indstype(arraytype, dims) +end + +@traitfn function similartype( + arraytype::Type{ArrayT}, ndims::NDims + ) where {{ArrayT; !IsWrappedArray{ArrayT}}} + return set_ndims(arraytype, ndims) +end + +function similartype( + arraytype::Type{<:AbstractArray}, dim1::Base.DimOrInd, dim_rest::Base.DimOrInd... + ) + return similartype(arraytype, (dim1, dim_rest...)) +end + +## Wrapped arrays +@traitfn function similartype(arraytype::Type{ArrayT}) where {{ArrayT; IsWrappedArray{ArrayT}}} + return similartype(unwrap_array_type(arraytype), NDims(arraytype)) +end + +@traitfn function similartype( + arraytype::Type{ArrayT}, eltype::Type + ) where {{ArrayT; IsWrappedArray{ArrayT}}} + return similartype(unwrap_array_type(arraytype), eltype, NDims(arraytype)) +end + +@traitfn function similartype( + arraytype::Type{ArrayT}, dims::Tuple + ) where {{ArrayT; IsWrappedArray{ArrayT}}} + return similartype(unwrap_array_type(arraytype), dims) +end + +@traitfn function similartype( + arraytype::Type{ArrayT}, ndims::NDims + ) where {{ArrayT; IsWrappedArray{ArrayT}}} + return similartype(unwrap_array_type(arraytype), ndims) +end + +# This is for uniform `Diag` storage which uses +# a Number as the data type. +# TODO: Delete this when we change to using a +# `FillArray` instead. This is a stand-in +# to make things work with the current design. +function similartype(numbertype::Type{<:Number}) + return numbertype +end + +# Instances +function similartype(array::AbstractArray, eltype::Type, dims...) + return similartype(typeof(array), eltype, dims...) +end +similartype(array::AbstractArray, eltype::Type) = similartype(typeof(array), eltype) +similartype(array::AbstractArray, dims...) = similartype(typeof(array), dims...) diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/ndims.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/ndims.jl new file mode 100644 index 0000000000..f640763583 --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/ndims.jl @@ -0,0 +1,6 @@ +struct NDims{ndims} end +Base.ndims(::NDims{ndims}) where {ndims} = ndims + +NDims(ndims::Integer) = NDims{ndims}() +NDims(arraytype::Type{<:AbstractArray}) = NDims(ndims(arraytype)) +NDims(array::AbstractArray) = NDims(typeof(array)) diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/type_parameters.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/type_parameters.jl new file mode 100644 index 0000000000..301ad20abf --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/type_parameters.jl @@ -0,0 +1,331 @@ +""" + struct Position{P} end + +Singleton type to statically represent the type-parameter position. +This is meant for internal use as a `Val`-like structure to improve type-inference. +""" +struct Position{P} end +Position(pos::Int) = Position{pos}() + +Base.Int(pos::Position) = Int(typeof(pos)) +Base.Int(::Type{Position{P}}) where {P} = Int(P) +Base.to_index(pos::Position) = Base.to_index(typeof(pos)) +Base.to_index(::Type{P}) where {P <: Position} = Int(P) + +""" + position(type::Type, position_name)::Position + +An optional interface function. Defining this allows accessing a parameter +at the defined position using the `position_name`. + +For example, defining `TypeParameterAccessors.position(::Type{<:MyType}, ::typeof(eltype)) = Position(1)` +allows accessing the first type parameter with `type_parameters(MyType(...), eltype)`, +in addition to the standard `type_parameters(MyType(...), 1)` or `type_parameters(MyType(...), Position(1))`. +""" +function position end + +position(object, name) = position(typeof(object), name) +position(::Type, pos::Int) = Position(pos) +position(::Type, pos::Position) = pos + +function position(type::Type, name) + type′ = unspecify_type_parameters(type) + if type === type′ + # Fallback definition that determines the + # position automatically from the supertype of + # the type. + return position_from_supertype(type′, name) + end + return position(type′, name) +end + +# Automatically determine the position of a type parameter of a type given +# a supertype and the name of the parameter. +function position_from_supertype(type::Type, name) + type′ = unspecify_type_parameters(type) + supertype_pos = position(supertype(type′), name) + return position_from_supertype_position(type′, supertype_pos) +end + +# Automatically determine the position of a type parameter of a type given +# the supertype and the position of the corresponding parameter in the supertype. +@generated function position_from_supertype_position( + ::Type{T}, supertype_pos::Position + ) where {T} + T′ = unspecify_type_parameters(T) + # The type parameters of the type as TypeVars. + # TODO: Ideally we would use `get_type_parameters` + # but that sometimes loses TypeVar names: + # https://github.com/ITensor/TypeParameterAccessors.jl/issues/30 + type_params = Base.unwrap_unionall(T′).parameters + # The type parameters of the immediate supertype as TypeVars. + # This has TypeVars with names that correspond to the names of + # the TypeVars of the type parameters of `T`, for example: + # ```julia + # julia> struct MyArray{B,A} <: AbstractArray{A,B} end + # + # julia> Base.unwrap_unionall(MyArray).parameters + # svec(B, A) + # + # julia> Base.unwrap_unionall(supertype(MyArray)).parameters + # svec(A, B) + # ``` + supertype_params = Base.unwrap_unionall(supertype(T′)).parameters + supertype_param = supertype_params[Int(supertype_pos)] + if !(supertype_param isa TypeVar) + error("Position not found.") + end + pos = findfirst(param -> (param.name == supertype_param.name), type_params) + if isnothing(pos) + return error("Position not found.") + end + return :(@inline; $(Position(pos))) +end + +function positions(::Type{T}, pos::Tuple) where {T} + return ntuple(length(pos)) do i + return position(T, pos[i]) + end +end + +""" + get_type_parameters(type_or_obj, [pos]) + +Return a tuple containing the type parameters of a given type or object. +Optionally you can specify a position to just get the parameter for that position, +or a tuple of positions to get a subset of parameters. + +If parameters are unspecified, returns a `TypeVar`. For a checked version, +see [`type_parameters`](@ref). +""" +function get_type_parameters end + +# This implementation is type-stable in 1.11, but not in 1.10. +# Attempts with `Base.@constprop :aggressive` failed, so generated function instead +# @inline type_parameters(::Type{T}) where {T} = Tuple(Base.unwrap_unionall(T).parameters) +@generated function get_type_parameters(::Type{T}) where {T} + params = wrap_symbol_quotenode.(Tuple(Base.unwrap_unionall(T).parameters)) + return :(@inline; ($(params...),)) +end +@inline get_type_parameters(::Type{T}, pos) where {T} = get_type_parameters( + T, position(T, pos) +) +@inline get_type_parameters(::Type{T}, ::Position{p}) where {T, p} = get_type_parameters(T)[p] +@inline get_type_parameters(::Type{T}, ::Position{0}) where {T} = T +@inline get_type_parameters(::Type{T}, pos::Tuple) where {T} = get_type_parameters.(T, pos) +@inline get_type_parameters(object, pos) = get_type_parameters(typeof(object), pos) +@inline get_type_parameters(object) = get_type_parameters(typeof(object)) + +""" + type_parameters(type_or_obj, [pos]) + +Return a tuple containing the type parameters of a given type or object. +Optionally you can specify a position to just get the parameter for that position, +or a tuple of positions to get a subset of parameters. + +Errors if parameters are unspecified. For an unchecked version, +see [`get_type_parameters`](@ref). +""" +function type_parameters end + +function type_parameters(::Type{T}) where {T} + params = get_type_parameters(T) + any(param -> param isa TypeVar, params) && + return error("One or more parameter is not specified.") + return params +end +@inline function type_parameters(::Type{T}, pos) where {T} + param = get_type_parameters(T, pos) + param isa TypeVar && return error("The parameter is not specified.") + return param +end +@inline type_parameters(::Type{T}, pos::Tuple) where {T} = type_parameters.(T, pos) +@inline type_parameters(object, pos) = type_parameters(typeof(object), pos) +@inline type_parameters(object) = type_parameters(typeof(object)) + +""" + nparameters(type_or_obj) + +Return the number of type parameters for a given type or object. +""" +nparameters(object) = nparameters(typeof(object)) +nparameters(::Type{T}) where {T} = length(get_type_parameters(T)) + +""" + is_parameter_specified(type::Type, pos) + +Return whether or not the type parameter at a given position is considered specified. +""" +function is_parameter_specified(::Type{T}, pos) where {T} + return !(get_type_parameters(T, pos) isa TypeVar) +end + +""" + unspecify_type_parameters(type::Type, [positions::Tuple]) + unspecify_type_parameters(type::Type, position) + +Return a new type where the type parameters at the given positions are unset. +""" +unspecify_type_parameters(::Type{T}) where {T} = Base.typename(T).wrapper +function unspecify_type_parameters(::Type{T}, pos::Tuple) where {T} + @inline + return unspecify_type_parameters(T, positions(T, pos)) +end +@generated function unspecify_type_parameters( + ::Type{T}, positions::Tuple{Vararg{Position}} + ) where {T} + allparams = collect(Any, get_type_parameters(T)) + for pos in type_parameters(positions) + allparams[pos] = get_type_parameters(unspecify_type_parameters(T), Int(pos)) + end + type_expr = construct_type_expr(T, allparams) + return :(@inline; $type_expr) +end +unspecify_type_parameters(::Type{T}, pos) where {T} = unspecify_type_parameters(T, (pos,)) + +""" + set_type_parameters(type::Type, positions::Tuple, parameters::Tuple) + set_type_parameters(type::Type, position, parameter) + +Return a new type where the type parameters at the given positions are set to the provided values. +""" +function set_type_parameters( + ::Type{T}, pos::Tuple{Vararg{Any, N}}, parameters::Tuple{Vararg{Any, N}} + ) where {T, N} + return set_type_parameters(T, positions(T, pos), parameters) +end +@generated function set_type_parameters( + ::Type{T}, positions::Tuple{Vararg{Position, N}}, params::Tuple{Vararg{Any, N}} + ) where {T, N} + # collect parameters and change + allparams = collect(Any, get_type_parameters(T)) + for (i, pos) in enumerate(get_type_parameters(positions)) + allparams[pos] = :(params[$i]) + end + type_expr = construct_type_expr(T, allparams) + return :(@inline; $type_expr) +end +function set_type_parameters(::Type{T}, pos, param) where {T} + return set_type_parameters(T, (pos,), (param,)) +end + +""" + specify_type_parameters(type::Type, positions::Tuple, parameters::Tuple) + specify_type_parameters(type::Type, position, parameter) + +Return a new type where the type parameters at the given positions are set to the provided values, +only if they were previously unspecified. +""" +function specify_type_parameters( + ::Type{T}, pos::Tuple{Vararg{Any, N}}, parameters::Tuple{Vararg{Any, N}} + ) where {T, N} + return specify_type_parameters(T, positions(T, pos), parameters) +end +function specify_type_parameters(::Type{T}, parameters::Tuple) where {T} + return specify_type_parameters(T, ntuple(identity, nparameters(T)), parameters) +end +@generated function specify_type_parameters( + ::Type{T}, positions::Tuple{Vararg{Position, N}}, params::Tuple{Vararg{Any, N}} + ) where {T, N} + # collect parameters and change unspecified + allparams = collect(Any, get_type_parameters(T)) + for (i, pos) in enumerate(get_type_parameters(positions)) + if !is_parameter_specified(T, pos()) + allparams[pos] = :(params[$i]) + end + end + type_expr = construct_type_expr(T, allparams) + return :(@inline; $type_expr) +end +function specify_type_parameters(::Type{T}, pos, param) where {T} + return specify_type_parameters(T, (pos,), (param,)) +end + +""" + default_type_parameters(type::Type)::Tuple + +An optional interface function. Defining this allows filling type parameters +of the specified type with default values. + +This function should output a Tuple of the default values, with exactly +one for each type parameter slot of the type. +""" +function default_type_parameters(::Type{T}, pos) where {T} + return default_type_parameters(T, position(T, pos)) +end +function default_type_parameters(::Type{T}, ::Position{pos}) where {T, pos} + param = default_type_parameters(T)[pos] + if param isa UndefinedDefaultTypeParameter + return error("No default parameter defined at this position.") + end + return param +end +default_type_parameters(::Type{T}, pos::Tuple) where {T} = default_type_parameters.(T, pos) +default_type_parameters(t, pos) = default_type_parameters(typeof(t), pos) +default_type_parameters(t) = default_type_parameters(typeof(t)) +function default_type_parameters(type::Type) + type′ = unspecify_type_parameters(type) + if type === type′ + return default_type_parameters_from_supertype(type′) + end + return default_type_parameters(type′) +end + +struct UndefinedDefaultTypeParameter end + +# Determine the default type parameters of a type from the default type +# parameters of the supertype of the type. Uses similar logic as +# `position_from_supertype_position` for matching TypeVar names +# between the type and the supertype. Type parameters that exist +# in the type but not the supertype will have a default type parameter +# `UndefinedDefaultTypeParameter()`. Accessing those type parameters +# by name/position will throw an error. +@generated function default_type_parameters_from_supertype(::Type{T}) where {T} + T′ = unspecify_type_parameters(T) + supertype_default_type_params = default_type_parameters(supertype(T′)) + type_params = Base.unwrap_unionall(T′).parameters + supertype_params = Base.unwrap_unionall(supertype(T′)).parameters + defaults = Any[UndefinedDefaultTypeParameter() for _ in 1:nparameters(T′)] + for (supertype_param, supertype_default_type_param) in + zip(supertype_params, supertype_default_type_params) + if !(supertype_param isa TypeVar) + continue + end + param_position = findfirst(param -> (param.name == supertype_param.name), type_params) + defaults[param_position] = supertype_default_type_param + end + return :(@inline; $(Tuple(defaults))) +end + +""" + set_default_type_parameters(type::Type, [positions::Tuple]) + set_default_type_parameters(type::Type, position) + +Set the type parameters at the given positions to their default values. +""" +function set_default_type_parameters(::Type{T}, pos::Tuple) where {T} + return set_type_parameters(T, pos, default_type_parameters.(T, pos)) +end +function set_default_type_parameters(::Type{T}) where {T} + return set_default_type_parameters(T, ntuple(identity, nparameters(T))) +end +function set_default_type_parameters(::Type{T}, pos) where {T} + return set_default_type_parameters(T, (pos,)) +end + +""" + specify_default_type_parameters(type::Type, [positions::Tuple]) + specify_default_type_parameters(type::Type, position) + +Set the type parameters at the given positions to their default values, if they +had not been specified. +""" +function specify_default_type_parameters(::Type{T}, pos::Tuple) where {T} + return specify_type_parameters(T, pos, default_type_parameters.(T, pos)) +end +function specify_default_type_parameters(::Type{T}) where {T} + return specify_default_type_parameters(T, ntuple(identity, nparameters(T))) +end +function specify_default_type_parameters(::Type{T}, pos) where {T} + return specify_default_type_parameters(T, (pos,)) +end diff --git a/NDTensors/src/vendored/TypeParameterAccessors/src/type_utils.jl b/NDTensors/src/vendored/TypeParameterAccessors/src/type_utils.jl new file mode 100644 index 0000000000..a0779d162d --- /dev/null +++ b/NDTensors/src/vendored/TypeParameterAccessors/src/type_utils.jl @@ -0,0 +1,14 @@ +"replace `Symbol`s with `QuoteNode`s to avoid expression interpolation" +wrap_symbol_quotenode(param) = param isa Symbol ? QuoteNode(param) : param + +"Construct the expression for qualifying a type with given parameters" +function construct_type_expr(type, parameters) + basetype = unspecify_type_parameters(type) + type_expr = Expr(:curly, basetype, wrap_symbol_quotenode.(parameters)...) + for parameter in reverse(parameters) + if parameter isa TypeVar + type_expr = Expr(:call, :UnionAll, parameter, type_expr) + end + end + return type_expr +end diff --git a/NDTensors/test/Project.toml b/NDTensors/test/Project.toml index 90742633a0..9193c2fefc 100644 --- a/NDTensors/test/Project.toml +++ b/NDTensors/test/Project.toml @@ -20,7 +20,6 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] From 09b468ae40a4918e2adf494aa791ce80b3562277 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sat, 11 Oct 2025 11:57:33 -0400 Subject: [PATCH 4/7] Add package extensions --- NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl | 2 +- .../ext/NDTensorsAMDGPUExt/linearalgebra.jl | 2 +- NDTensors/ext/NDTensorsAMDGPUExt/set_types.jl | 2 +- .../ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl | 2 ++ NDTensors/ext/NDTensorsCUDAExt/adapt.jl | 2 +- .../ext/NDTensorsCUDAExt/linearalgebra.jl | 2 +- NDTensors/ext/NDTensorsCUDAExt/set_types.jl | 2 +- .../ext/NDTensorsGPUArraysCoreExt/contract.jl | 2 +- .../NDTensorsJLArraysExt.jl | 2 ++ .../ext/NDTensorsJLArraysExt/linearalgebra.jl | 2 +- NDTensors/ext/NDTensorsMetalExt/adapt.jl | 2 +- .../ext/NDTensorsMetalExt/linearalgebra.jl | 2 +- NDTensors/ext/NDTensorsMetalExt/set_types.jl | 2 +- NDTensors/src/NDTensors.jl | 18 +++++++++++++++++- .../AMDGPUExtensions/src/AMDGPUExtensions.jl | 9 --------- NDTensors/src/lib/AMDGPUExtensions/src/roc.jl | 2 +- .../lib/CUDAExtensions/src/CUDAExtensions.jl | 9 --------- NDTensors/src/lib/CUDAExtensions/src/cuda.jl | 2 +- NDTensors/src/lib/Expose/src/Expose.jl | 9 --------- NDTensors/src/lib/Expose/src/exposed.jl | 2 +- .../src/GPUArraysCoreExtensions.jl | 9 --------- .../src/gpuarrayscore.jl | 2 +- .../lib/MetalExtensions/src/MetalExtensions.jl | 9 --------- NDTensors/src/lib/MetalExtensions/src/metal.jl | 2 +- .../RankFactorization/src/RankFactorization.jl | 9 --------- .../RankFactorization/src/default_kwargs.jl | 2 +- .../RankFactorization/src/truncate_spectrum.jl | 2 +- .../ext/TypeParameterAccessorsAMDGPUExt.jl | 12 ++++++++++++ .../ext/TypeParameterAccessorsCUDAExt.jl | 12 ++++++++++++ .../ext/TypeParameterAccessorsFillArraysExt.jl | 17 +++++++++++++++++ .../ext/TypeParameterAccessorsJLArraysExt.jl | 10 ++++++++++ .../ext/TypeParameterAccessorsMetalExt.jl | 12 ++++++++++++ .../TypeParameterAccessorsStridedViewsExt.jl | 15 +++++++++++++++ .../ext/TypeParameterAccessorsoneAPIExt.jl | 12 ++++++++++++ NDTensors/vendor/run.jl | 2 ++ 35 files changed, 131 insertions(+), 73 deletions(-) create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsAMDGPUExt.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsCUDAExt.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsFillArraysExt.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsJLArraysExt.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsMetalExt.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsStridedViewsExt.jl create mode 100644 NDTensors/src/vendored/TypeParameterAccessors/ext/TypeParameterAccessorsoneAPIExt.jl diff --git a/NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl b/NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl index 520c4dabad..6106af83d1 100644 --- 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 index fcc9edcd36..8433db5987 100644 --- 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 index 8bebb320d2..508804bf61 100644 --- 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 index 2a7a458ff9..e3540529fb 100644 --- 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 index f6a97ac593..0009ccde67 100644 --- 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 index f850da1d39..4c5d1bfd2a 100644 --- 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 index 21205aec11..d613ee1dca 100644 --- 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 index d87e185264..3fc4d80437 100644 --- 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 index 8de2f1dcfd..7e736fb838 100644 --- 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 index b8079ec973..12fb3281f0 100644 --- 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 index 217b2e697e..a7d31fbd2e 100644 --- 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 index 9c11b51eeb..b05e6d7d62 100644 --- 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 index 34c331d1b2..ea44e4a7e8 100644 --- 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 index fbb1a320bd..7a30f85781 100644 --- 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 index 2007c8f1ef..cc6eb142a8 100644 --- 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 index 0aa1a58693..d3bc16d4e9 100644 --- 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 index e9ae01be1b..b2ff7ed626 100644 --- 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 index 60dad8d5d3..b5a5b4cf5a 100644 --- 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 index 297fffd4ae..470bdfe39b 100644 --- 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 index 3ebc88e416..648a7f7e42 100644 --- 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 index 9f6972403b..4ea8d95a1f 100644 --- 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 index 1e37904042..7271023200 100644 --- 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 index deadddfbdf..9bfd6871f9 100644 --- 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 index 3253c7ed66..65d1dcffb7 100644 --- 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 index d41bbbc317..f5a92a8df9 100644 --- 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 index 9cb170b0e5..2cb6edbaa4 100644 --- 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 index d05c2b2ebc..9ad315efc6 100644 --- 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 new file mode 100644 index 0000000000..e4112b31bb --- /dev/null +++ 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 new file mode 100644 index 0000000000..20534e6542 --- /dev/null +++ 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 new file mode 100644 index 0000000000..28db107ab4 --- /dev/null +++ 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 new file mode 100644 index 0000000000..297ab458c6 --- /dev/null +++ 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 new file mode 100644 index 0000000000..e55cd5d883 --- /dev/null +++ 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 new file mode 100644 index 0000000000..09daaea787 --- /dev/null +++ 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 new file mode 100644 index 0000000000..899aebe8d2 --- /dev/null +++ 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 index 2491b21f5a..716afc8fc8 100644 --- 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 From 183786d57d85d5db5707a1551ee19a0e67b55cb6 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sat, 11 Oct 2025 12:23:51 -0400 Subject: [PATCH 5/7] Fix GPU backends --- NDTensors/ext/NDTensorsAMDGPUExt/NDTensorsAMDGPUExt.jl | 9 +++++++++ NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl | 9 +++++++++ .../ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl | 9 +++++++++ NDTensors/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl | 9 +++++++++ 4 files changed, 36 insertions(+) diff --git a/NDTensors/ext/NDTensorsAMDGPUExt/NDTensorsAMDGPUExt.jl b/NDTensors/ext/NDTensorsAMDGPUExt/NDTensorsAMDGPUExt.jl index 34004a7b7f..e304c2038f 100644 --- a/NDTensors/ext/NDTensorsAMDGPUExt/NDTensorsAMDGPUExt.jl +++ b/NDTensors/ext/NDTensorsAMDGPUExt/NDTensorsAMDGPUExt.jl @@ -1,5 +1,14 @@ module NDTensorsAMDGPUExt +module Vendored + include( + joinpath( + "..", "..", "src", "vendored", "TypeParameterAccessors", "ext", + "TypeParameterAccessorsAMDGPUExt.jl" + ) + ) +end + include("append.jl") include("copyto.jl") include("set_types.jl") diff --git a/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl b/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl index e3540529fb..716d57efd8 100644 --- a/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl +++ b/NDTensors/ext/NDTensorsCUDAExt/NDTensorsCUDAExt.jl @@ -1,5 +1,14 @@ module NDTensorsCUDAExt +module Vendored + include( + joinpath( + "..", "..", "src", "vendored", "TypeParameterAccessors", "ext", + "TypeParameterAccessorsCUDAExt.jl" + ) + ) +end + include("append.jl") include("default_kwargs.jl") include("copyto.jl") diff --git a/NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl b/NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl index 7e736fb838..396ee2c829 100644 --- a/NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl +++ b/NDTensors/ext/NDTensorsJLArraysExt/NDTensorsJLArraysExt.jl @@ -1,5 +1,14 @@ module NDTensorsJLArraysExt +module Vendored + include( + joinpath( + "..", "..", "src", "vendored", "TypeParameterAccessors", "ext", + "TypeParameterAccessorsJLArraysExt.jl" + ) + ) +end + include("copyto.jl") include("indexing.jl") include("linearalgebra.jl") diff --git a/NDTensors/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl b/NDTensors/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl index 8cca86ca2a..b7ef39aadd 100644 --- a/NDTensors/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl +++ b/NDTensors/ext/NDTensorsMetalExt/NDTensorsMetalExt.jl @@ -1,5 +1,14 @@ module NDTensorsMetalExt +module Vendored + include( + joinpath( + "..", "..", "src", "vendored", "TypeParameterAccessors", "ext", + "TypeParameterAccessorsMetalExt.jl" + ) + ) +end + include("adapt.jl") include("set_types.jl") include("indexing.jl") From d87df93f0bd01516108e48befbe0a6a66075debf Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sat, 11 Oct 2025 12:38:54 -0400 Subject: [PATCH 6/7] Update vendor script to fix using statements --- NDTensors/vendor/run.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/NDTensors/vendor/run.jl b/NDTensors/vendor/run.jl index 716afc8fc8..33311db5b0 100644 --- a/NDTensors/vendor/run.jl +++ b/NDTensors/vendor/run.jl @@ -15,4 +15,12 @@ for (name, version) in deps cp(joinpath(local_path, "src"), joinpath(p, "src"); force = true) mkpath(joinpath(p, "ext")) cp(joinpath(local_path, "ext"), joinpath(p, "ext"); force = true) + + for filename in readdir(joinpath(p, "ext"); join = true) + txt = read(filename, String) + open(filename, "w") do f + replacement = "using $name" => "using NDTensors.Vendored.$name" + return write(f, replace(txt, replacement)) + end + end end From fd291a16b949a9bc70ad04d12e48b7aa57a4e679 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sat, 11 Oct 2025 12:43:02 -0400 Subject: [PATCH 7/7] Update script, credit ExplicitImports.jl --- NDTensors/vendor/README.md | 2 ++ NDTensors/vendor/run.jl | 2 ++ 2 files changed, 4 insertions(+) diff --git a/NDTensors/vendor/README.md b/NDTensors/vendor/README.md index ebd43458f0..604655663c 100644 --- a/NDTensors/vendor/README.md +++ b/NDTensors/vendor/README.md @@ -11,3 +11,5 @@ To update the dependencies, install the `vendor` environment and invoke `run.jl` julia --project=vendor -e 'using Pkg; Pkg.instantiate()' julia --project=vendor vendor/run.jl ``` + +Note that this is based on https://github.com/JuliaTesting/ExplicitImports.jl/tree/v1.13.2/vendor. diff --git a/NDTensors/vendor/run.jl b/NDTensors/vendor/run.jl index 33311db5b0..19bcbaf118 100644 --- a/NDTensors/vendor/run.jl +++ b/NDTensors/vendor/run.jl @@ -1,3 +1,5 @@ +# Based on https://github.com/JuliaTesting/ExplicitImports.jl/tree/v1.13.2/vendor + using PackageAnalyzer: PackageAnalyzer deps = [("TypeParameterAccessors", v"0.3.11")]