Updates to similartype

ITensor · Feb 28, 2024 · 591e547 · 591e547
1 parent 237beff
commit 591e547
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 18 deletions.
diff --git a/NDTensors/src/abstractarray/ndims.jl b/NDTensors/src/abstractarray/ndims.jl
@@ -5,7 +5,7 @@ ndims((array::AbstractArray)) = ndims(typeof(array))
 ndims(arraytype::Type{<:AbstractArray}) = type_parameter(arraytype, Base.ndims)
 
 ## TODO for now have `NDTensors.set_ndims` call `TypeParameterAccessors.set_ndims`
-set_ndims(type::Type, length) = TypeParameterAccessors.set_ndims(type, length)
+NDTensors.set_ndims(type::Type, length) = TypeParameterAccessors.set_ndims(type, length)
 
 TypeParameterAccessors.position(::Type{<:Number}, ::typeof(ndims)) = Self()
 

diff --git a/NDTensors/src/abstractarray/similar.jl b/NDTensors/src/abstractarray/similar.jl
@@ -1,6 +1,5 @@
 ## todo working on this still
-using .TypeParameterAccessors: set_eltype
-using .Unwrap: IsWrappedArray
+using .TypeParameterAccessors: IsWrappedArray, set_eltype, specify_default_parameters, unwrap_array_type
 
 ## Custom `NDTensors.similar` implementation.
 ## More extensive than `Base.similar`.
@@ -16,6 +15,8 @@ end
 # Catches conversions of dimensions specified by ranges
 # dimensions specified by integers with `Base.to_shape`.
 # NDTensors.similar
+## TODO this fails when types are partially specified. Should we modify the
+## behavior?
 function similar(arraytype::Type{<:AbstractArray}, dims::Dims)
   return similartype(arraytype, dims)(undef, dims)
 end
@@ -105,18 +106,13 @@ function similartype(arraytype::Type{<:AbstractArray}, eltype::Type, dims::Tuple
   return similartype(similartype(arraytype, eltype), dims)
 end
 
-@traitfn function similartype(
-  arraytype::Type{ArrayT}, eltype::Type
-) where {ArrayT; !IsWrappedArray{ArrayT}}
+## Set eltype captures WrappedArray types as long as `position(::Type, ::typeof(parenttype))` is defined
+function similartype(
+  arraytype::Type{<:AbstractArray}, eltype::Type
+)
   return set_eltype(arraytype, eltype)
 end
 
-@traitfn function similartype(
-  arraytype::Type{ArrayT}, dims::Tuple
-) where {ArrayT; !IsWrappedArray{ArrayT}}
-  return set_indstype(arraytype, dims)
-end
-
 function similartype(arraytype::Type{<:AbstractArray}, dims::DimOrInd...)
   return similartype(arraytype, dims)
 end
@@ -125,17 +121,16 @@ function similartype(arraytype::Type{<:AbstractArray})
   return similartype(arraytype, eltype(arraytype))
 end
 
-## Wrapped arrays
 @traitfn function similartype(
-  arraytype::Type{ArrayT}, eltype::Type
-) where {ArrayT; IsWrappedArray{ArrayT}}
-  return similartype(unwrap_type(arraytype), eltype)
+  arraytype::Type{ArrayT}, dims::Tuple
+) where {ArrayT; !IsWrappedArray{ArrayT}}
+  return set_indstype(arraytype, dims)
 end
 
 @traitfn function similartype(
   arraytype::Type{ArrayT}, dims::Tuple
 ) where {ArrayT; IsWrappedArray{ArrayT}}
-  return similartype(unwrap_type(arraytype), dims)
+  return set_parenttype(arraytype, similartype(unwrap_array_type(arraytype), dims))
 end
 
 # This is for uniform `Diag` storage which uses

diff --git a/NDTensors/src/lib/TypeParameterAccessors/src/default_parameters.jl b/NDTensors/src/lib/TypeParameterAccessors/src/default_parameters.jl
@@ -1,4 +1,4 @@
-default_type_parameters(type::Type) = error("Not implemented")
+default_type_parameters(type::Type) = error("The default parameters of $(type) is not defined. Please specify the parameters manually.")
 function default_type_parameters(type::Type, positions::Tuple)
   return default_type_parameter.(type, positions)
 end