[NDTensors] [ITensors] Excise unneeded submodules

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.74"
+version = "0.4.0"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -29,6 +29,7 @@ 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]
@@ -89,6 +90,7 @@ StridedViews = "0.2.2, 0.3"
 TBLIS = "0.2"
 TimerOutputs = "0.5.5"
 TupleTools = "1.2.0"
+TypeParameterAccessors = "0.2"
 VectorInterface = "0.4.2, 0.5"
 cuTENSOR = "2"
 julia = "1.10"

NDTensors/ext/NDTensorsAMDGPUExt/adapt.jl

@@ -2,16 +2,14 @@ using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
 using NDTensors.AMDGPUExtensions: AMDGPUExtensions, ROCArrayAdaptor
 using NDTensors.GPUArraysCoreExtensions: storagemode
 using NDTensors.TypeParameterAccessors:
-  default_type_parameter,
-  set_type_parameter,
-  set_type_parameters,
-  type_parameter,
-  type_parameters
+  default_type_parameters, set_type_parameters, type_parameters
 using Adapt: Adapt, adapt
 using AMDGPU: AMDGPU, ROCArray, ROCVector
 using Functors: fmap
 
-function AMDGPUExtensions.roc(xs; storagemode=default_type_parameter(ROCArray, storagemode))
+function AMDGPUExtensions.roc(
+  xs; storagemode=default_type_parameters(ROCArray, storagemode)
+)
   return fmap(x -> adapt(ROCArrayAdaptor{storagemode}(), x), xs)
 end
 

NDTensors/ext/NDTensorsCUDAExt/adapt.jl

@@ -5,9 +5,9 @@ using NDTensors: NDTensors, EmptyStorage, adapt_storagetype, emptytype
 using NDTensors.CUDAExtensions: CUDAExtensions, CuArrayAdaptor
 using NDTensors.GPUArraysCoreExtensions: storagemode
 using NDTensors.TypeParameterAccessors:
-  default_type_parameter, set_type_parameters, type_parameters
+  default_type_parameters, set_type_parameters, type_parameters
 
-function CUDAExtensions.cu(xs; storagemode=default_type_parameter(CuArray, storagemode))
+function CUDAExtensions.cu(xs; storagemode=default_type_parameters(CuArray, storagemode))
   return fmap(x -> adapt(CuArrayAdaptor{storagemode}(), x), xs)
 end
 

NDTensors/ext/NDTensorsGPUArraysCoreExt/blocksparsetensor.jl

@@ -1,6 +1,5 @@
 using GPUArraysCore: @allowscalar, AbstractGPUArray
-using NDTensors: NDTensors, BlockSparseTensor, dense, diag, map_diag!
-using NDTensors.DiagonalArrays: diaglength
+using NDTensors: NDTensors, BlockSparseTensor, dense, diag, diaglength, map_diag!
 using NDTensors.Expose: Exposed, unexpose
 
 ## TODO to circumvent issues with blocksparse and scalar indexing
@@ -11,7 +10,7 @@ function NDTensors.diag(ETensor::Exposed{<:AbstractGPUArray,<:BlockSparseTensor}
   return diag(dense(unexpose(ETensor)))
 end
 
-## TODO scalar indexing is slow here 
+## TODO scalar indexing is slow here
 function NDTensors.map_diag!(
   f::Function,
   exposed_t_destination::Exposed{<:AbstractGPUArray,<:BlockSparseTensor},

NDTensors/src/NDTensors.jl

@@ -17,6 +17,7 @@ include("abstractarray/similar.jl")
 include("abstractarray/mul.jl")
 include("abstractarray/permutedims.jl")
 include("abstractarray/generic_array_constructors.jl")
+include("abstractarray/diaginterface.jl")
 include("array/permutedims.jl")
 include("array/mul.jl")
 include("tupletools.jl")
@@ -91,15 +92,6 @@ include("empty/adapt.jl")
 #
 include("deprecated.jl")
 
-#####################################
-# NDTensorsNamedDimsArraysExt
-# I tried putting this inside of an
-# `NDTensorsNamedDimsArraysExt` module
-# but for some reason it kept overloading
-# `Base.similar` instead of `NDTensors.similar`.
-#
-include("NDTensorsNamedDimsArraysExt/NDTensorsNamedDimsArraysExt.jl")
-
 #####################################
 # A global timer used with TimerOutputs.jl
 #

NDTensors/src/abstractarray/diaginterface.jl

@@ -0,0 +1,29 @@
+# Selected interface functions from https://github.com/ITensor/DiagonalArrays.jl,
+# copied here so we don't have to depend on `DiagonalArrays.jl`.
+
+function diaglength(a::AbstractArray)
+  return minimum(size(a))
+end
+
+function diagstride(a::AbstractArray)
+  s = 1
+  p = 1
+  for i in 1:(ndims(a) - 1)
+    p *= size(a, i)
+    s += p
+  end
+  return s
+end
+
+function diagindices(a::AbstractArray)
+  maxdiag = LinearIndices(a)[CartesianIndex(ntuple(Returns(diaglength(a)), ndims(a)))]
+  return 1:diagstride(a):maxdiag
+end
+
+function diagindices(a::AbstractArray{<:Any,0})
+  return Base.OneTo(1)
+end
+
+function diagview(a::AbstractArray)
+  return @view a[diagindices(a)]
+end

NDTensors/src/abstractarray/generic_array_constructors.jl

@@ -1,5 +1,8 @@
-using .TypeParameterAccessors:
-  unwrap_array_type, specify_default_type_parameters, type_parameter
+using TypeParameterAccessors:
+  unwrap_array_type,
+  specify_default_type_parameters,
+  specify_type_parameters,
+  type_parameters
 
 # Convert to Array, avoiding copying if possible
 array(a::AbstractArray) = a
@@ -8,9 +11,9 @@ vector(a::AbstractVector) = a
 
 ## Warning to use these functions it is necessary to define `TypeParameterAccessors.position(::Type{<:YourArrayType}, ::typeof(ndims)))`
 # Implementation, catches if `ndims(arraytype) != length(dims)`.
-## TODO convert ndims to `type_parameter(::, typeof(ndims))`
+## TODO convert ndims to `type_parameters(::, typeof(ndims))`
 function generic_randn(arraytype::Type{<:AbstractArray}, dims...; rng=Random.default_rng())
-  arraytype_specified = specify_type_parameter(
+  arraytype_specified = specify_type_parameters(
     unwrap_array_type(arraytype), ndims, length(dims)
   )
   arraytype_specified = specify_default_type_parameters(arraytype_specified)
@@ -27,7 +30,7 @@ end
 
 # Implementation, catches if `ndims(arraytype) != length(dims)`.
 function generic_zeros(arraytype::Type{<:AbstractArray}, dims...)
-  arraytype_specified = specify_type_parameter(
+  arraytype_specified = specify_type_parameters(
     unwrap_array_type(arraytype), ndims, length(dims)
   )
   arraytype_specified = specify_default_type_parameters(arraytype_specified)

NDTensors/src/abstractarray/iscu.jl

@@ -1,4 +1,4 @@
-using .TypeParameterAccessors: unwrap_array_type
+using TypeParameterAccessors: unwrap_array_type
 # TODO: Make `isgpu`, `ismtl`, etc.
 # For `isgpu`, will require a `NDTensorsGPUArrayCoreExt`.
 iscu(A::AbstractArray) = iscu(typeof(A))

NDTensors/src/abstractarray/set_types.jl

@@ -1,4 +1,4 @@
-using .TypeParameterAccessors: TypeParameterAccessors
+using TypeParameterAccessors: TypeParameterAccessors
 
 """
 # Do we still want to define things like this?

NDTensors/src/abstractarray/similar.jl

@@ -1,5 +1,5 @@
 using Base: DimOrInd, Dims, OneTo
-using .TypeParameterAccessors: IsWrappedArray, unwrap_array_type, set_eltype, similartype
+using TypeParameterAccessors: IsWrappedArray, unwrap_array_type, set_eltype, similartype
 
 ## Custom `NDTensors.similar` implementation.
 ## More extensive than `Base.similar`.

NDTensors/src/adapt.jl

@@ -27,11 +27,11 @@ double_precision(x) = fmap(x -> adapt(double_precision(eltype(x)), x), x)
 # Used to adapt `EmptyStorage` types
 #
 
-using .TypeParameterAccessors: specify_type_parameter, specify_type_parameters
+using TypeParameterAccessors: specify_type_parameters
 function adapt_storagetype(to::Type{<:AbstractVector}, x::Type{<:TensorStorage})
-  return set_datatype(x, specify_type_parameter(to, eltype, eltype(x)))
+  return set_datatype(x, specify_type_parameters(to, eltype, eltype(x)))
 end
 
 function adapt_storagetype(to::Type{<:AbstractArray}, x::Type{<:TensorStorage})
-  return set_datatype(x, specify_type_parameter(to, (ndims, eltype), (1, eltype(x))))
+  return set_datatype(x, specify_type_parameters(to, (ndims, eltype), (1, eltype(x))))
 end