Merge 2f38cef into b04a973

Author: kmp5VT

NDTensors/Project.toml

@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
@@ -49,3 +50,4 @@ Strided = "0.3, 1, 2"
 TimerOutputs = "0.5.5"
 TupleTools = "1.2.0"
 julia = "1.6"
+

NDTensors/ext/NDTensorsCUDAExt/adapt.jl

@@ -3,7 +3,7 @@
 ## Will not allow us to properly utilize the buffer preference without changing the value of
 ## default_buffertype. Also `adapt(CuVector{<:Any, <:Any, Buffertype})` fails to work properly
 struct NDTensorCuArrayAdaptor{B} end
-## TODO make this work for unified. This works but overwrites CUDA's adapt_storage. This fails for emptystorage...
+## TODO make this work for unified. This works but overwrites CUDA's adapt_storage.
 function cu(xs; unified::Bool=false)
   return fmap(
     x -> adapt(NDTensorCuArrayAdaptor{unified ? Mem.UnifiedBuffer : Mem.DeviceBuffer}(), x),
@@ -18,10 +18,3 @@ function Adapt.adapt_storage(adaptor::NDTensorCuArrayAdaptor, xs::AbstractArray)
   BufT = buffertype(adaptor)
   return isbits(xs) ? xs : CuArray{ElT,1,BufT}(xs)
 end
-
-function NDTensors.adapt_storagetype(
-  adaptor::NDTensorCuArrayAdaptor, xs::Type{EmptyStorage{ElT,StoreT}}
-) where {ElT,StoreT}
-  BufT = buffertype(adaptor)
-  return NDTensors.emptytype(NDTensors.adapt_storagetype(CuVector{ElT,BufT}, StoreT))
-end

NDTensors/ext/NDTensorsCUDAExt/imports.jl

@@ -1,4 +1,4 @@
-import NDTensors: cu, set_ndims, set_eltype, set_eltype_if_unspecified, similartype
+import NDTensors: cu, set_ndims, set_eltype, specify_eltype, similartype
 import NDTensors:
   ContractionProperties, _contract!, GemmBackend, auto_select_backend, _gemm!
 import NDTensors.SetParameters: nparameters, get_parameter, set_parameter, default_parameter

NDTensors/ext/NDTensorsMetalExt/imports.jl

@@ -1,4 +1,4 @@
-import NDTensors: mtl, set_ndims, set_eltype, set_eltype_if_unspecified
+import NDTensors: mtl, set_ndims, set_eltype, specify_eltype
 import NDTensors.SetParameters: nparameters, get_parameter, set_parameter, default_parameter
 
 using Metal: DefaultStorageMode

NDTensors/src/NDTensors.jl

@@ -4,6 +4,7 @@ using Adapt
 using Base.Threads
 using Compat
 using Dictionaries
+using FillArrays
 using FLoops
 using Folds
 using Random
@@ -32,6 +33,11 @@ using Base.Threads: @spawn
 include("imports.jl")
 include("exports.jl")
 
+#####################################
+# Unspecified Element Types
+#
+include("unspec_eltype/unspecified_zero.jl")
+
 #####################################
 # General functionality
 #
@@ -44,7 +50,6 @@ include("abstractarray/ndims.jl")
 include("abstractarray/fill.jl")
 include("array/set_types.jl")
 include("tupletools.jl")
-include("emptynumber.jl")
 include("nodata.jl")
 include("tensorstorage/tensorstorage.jl")
 include("tensorstorage/set_types.jl")
@@ -55,9 +60,18 @@ include("dims.jl")
 include("tensor/set_types.jl")
 include("tensor/similar.jl")
 include("adapt.jl")
+include("abstractarray/data_isa.jl")
 include("tensoralgebra/generic_tensor_operations.jl")
 include("tensoralgebra/contraction_logic.jl")
 
+#####################################
+# Zeros
+#
+include("zeros/unallocated_zeros.jl")
+include("zeros/similar.jl")
+include("zeros/set_types.jl")
+include("zeros/allocate.jl")
+
 #####################################
 # DenseTensor and DiagTensor
 #
@@ -101,14 +115,6 @@ include("blocksparse/similar.jl")
 include("blocksparse/combiner.jl")
 include("blocksparse/linearalgebra.jl")
 
-#####################################
-# Empty
-#
-include("empty/empty.jl")
-include("empty/EmptyTensor.jl")
-include("empty/tensoralgebra/contract.jl")
-include("empty/adapt.jl")
-
 #####################################
 # Deprecations
 #

NDTensors/src/abstractarray/data_isa.jl

@@ -0,0 +1,6 @@
+data_isa(t::Tensor, datatype::Type) = data_isa(storage(t), datatype)
+data_isa(s::TensorStorage, datatype::Type) = data_isa(data(s), datatype)
+data_isa(d::AbstractArray, datatype::Type) = d isa datatype # Might have to unwrap if it is reshaped, sliced, etc.
+## This is for the UniformDiag type can be removed
+## after switching to FillArrays.Diag
+data_isa(n::Number, datatype::Type) = false

NDTensors/src/abstractarray/fill.jl

@@ -10,6 +10,7 @@ function generic_randn(arraytype::Type{<:AbstractArray}, dim::Integer=0)
   return data
 end
 
+## TODO: Should generic Zeros by default construct an NDTensors.Zeros with an alloc type of `NDTensors.default_datatype()` ?
 function generic_zeros(arraytype::Type{<:AbstractArray}, dim::Integer=0)
   arraytype_specified = set_unspecified_parameters(
     leaf_parenttype(arraytype), DefaultParameters()

NDTensors/src/abstractarray/set_types.jl

@@ -37,30 +37,37 @@ function set_indstype(arraytype::Type{<:AbstractArray}, dims::Tuple)
   return set_ndims(arraytype, length(dims))
 end
 
-function set_eltype_if_unspecified(
+function specify_eltype(
   arraytype::Type{<:AbstractArray{T}}, eltype::Type=default_eltype()
 ) where {T}
   return arraytype
 end
 
 #TODO transition to set_eltype when working for wrapped types
-function set_eltype_if_unspecified(
-  arraytype::Type{<:AbstractArray}, eltype::Type=default_eltype()
-)
+function specify_eltype(arraytype::Type{<:AbstractArray}, eltype::Type=default_eltype())
   return similartype(arraytype, eltype)
 end
 
-function set_parameter_if_unspecified(
+const RealOrComplex{T} = Union{T,Complex{T}}
+function specify_eltype(
+  arraytype::Type{<:AbstractArray{<:RealOrComplex{UnspecifiedZero}}},
+  elt::Type=default_eltype(),
+)
+  elt = promote_type(elt, eltype(arraytype))
+  return set_eltype(arraytype, elt)
+end
+
+function specify_parameters(
   arraytype::Type{<:AbstractArray{ElT,N}}, eltype::Type=default_eltype(), ndims::Integer=1
 ) where {ElT,N}
   return arraytype
 end
-function set_parameter_if_unspecified(
+function specify_parameters(
   arraytype::Type{<:AbstractArray{ElT}}, eltype::Type=default_eltype(), ndims::Integer=1
 ) where {ElT}
   return set_ndims(arraytype, ndims)
 end
-function set_parameter_if_unspecified(
+function specify_parameters(
   arraytype::Type{<:AbstractArray}, eltype::Type=default_eltype(), ndims::Integer=1
 )
   return set_eltype(set_ndims(arraytype, ndims), eltype)

NDTensors/src/adapt.jl

@@ -4,7 +4,7 @@ adapt_structure(to, x::Tensor) = setstorage(x, adapt(to, storage(x)))
 cpu(eltype::Type{<:Number}, x) = fmap(x -> adapt(Array{eltype}, x), x)
 cpu(x) = fmap(x -> adapt(Array, x), x)
 
-# Implemented in `ITensorGPU` and NDTensorCUDA
+# Implemented in `ITensorGPU` and NDTensorsCUDAExt
 function cu end
 
 function mtl end
@@ -30,9 +30,9 @@ double_precision(x) = fmap(x -> adapt(double_precision(eltype(x)), x), x)
 #
 
 function adapt_storagetype(to::Type{<:AbstractVector}, x::Type{<:TensorStorage})
-  return set_datatype(x, set_eltype_if_unspecified(to, eltype(x)))
+  return set_datatype(x, specify_eltype(to, eltype(x)))
 end
 
 function adapt_storagetype(to::Type{<:AbstractArray}, x::Type{<:TensorStorage})
-  return set_datatype(x, set_eltype_if_unspecified(set_ndims(to, 1), eltype(x)))
+  return set_datatype(x, specify_eltype(set_ndims(to, 1), eltype(x)))
 end

NDTensors/src/blocksparse/blocksparse.jl

@@ -2,6 +2,7 @@
 # BlockSparse storage
 #
 
+## BlockSparse constructors
 struct BlockSparse{ElT,VecT,N} <: TensorStorage{ElT}
   data::VecT
   blockoffsets::BlockOffsets{N}  # Block number-offset pairs
@@ -12,36 +13,16 @@ struct BlockSparse{ElT,VecT,N} <: TensorStorage{ElT}
   end
 end
 
-# TODO: Implement as `fieldtype(storagetype, :data)`.
-datatype(::Type{<:BlockSparse{<:Any,DataT}}) where {DataT} = DataT
-# TODO: Implement as `ndims(blockoffsetstype(storagetype))`.
-ndims(storagetype::Type{<:BlockSparse{<:Any,<:Any,N}}) where {N} = N
-# TODO: Implement as `fieldtype(storagetype, :blockoffsets)`.
-blockoffsetstype(storagetype::Type{<:BlockSparse}) = BlockOffsets{ndims(storagetype)}
-
-function set_datatype(storagetype::Type{<:BlockSparse}, datatype::Type{<:AbstractVector})
-  return BlockSparse{eltype(datatype),datatype,ndims(storagetype)}
-end
-
-function set_ndims(storagetype::Type{<:BlockSparse}, ndims)
-  return BlockSparse{eltype(storagetype),datatype(storagetype),ndims}
-end
-
-# TODO: Write as `(::Type{<:BlockSparse})()`.
-BlockSparse{ElT,DataT,N}() where {ElT,DataT,N} = BlockSparse(DataT(), BlockOffsets{N}())
-
 function BlockSparse(
   datatype::Type{<:AbstractArray}, blockoffsets::BlockOffsets, dim::Integer; vargs...
 )
-  return BlockSparse(
-    fill!(NDTensors.similar(datatype, dim), zero(eltype(datatype))), blockoffsets; vargs...
-  )
+  return BlockSparse(generic_zeros(datatype, dim), blockoffsets; vargs...)
 end
 
 function BlockSparse(
   eltype::Type{<:Number}, blockoffsets::BlockOffsets, dim::Integer; vargs...
 )
-  return BlockSparse(Vector{eltype}, blockoffsets, dim; vargs...)
+  return BlockSparse(default_datatype(eltype), blockoffsets, dim; vargs...)
 end
 
 function BlockSparse(x::Number, blockoffsets::BlockOffsets, dim::Integer; vargs...)
@@ -51,16 +32,38 @@ end
 function BlockSparse(
   ::Type{ElT}, ::UndefInitializer, blockoffsets::BlockOffsets, dim::Integer; vargs...
 ) where {ElT<:Number}
-  return BlockSparse(Vector{ElT}(undef, dim), blockoffsets; vargs...)
+  return BlockSparse(default_datatype(ElT)(undef, dim), blockoffsets; vargs...)
+end
+
+function BlockSparse(::UndefInitializer, blockoffsets::BlockOffsets, dim::Integer; vargs...)
+  return BlockSparse(NDTensors.default_eltype(), undef, blockoffsets, dim; vargs...)
 end
 
 function BlockSparse(blockoffsets::BlockOffsets, dim::Integer; vargs...)
-  return BlockSparse(Float64, blockoffsets, dim; vargs...)
+  return BlockSparse(default_datatype(default_eltype()), blockoffsets, dim; vargs...)
 end
 
-function BlockSparse(::UndefInitializer, blockoffsets::BlockOffsets, dim::Integer; vargs...)
-  return BlockSparse(Float64, undef, blockoffsets, dim; vargs...)
+# TODO: Write as `(::Type{<:BlockSparse})()`.
+BlockSparse{ElT,DataT,N}() where {ElT,DataT,N} = BlockSparse(DataT(), BlockOffsets{N}())
+
+## End BlockSparse constructors
+
+## Blocksparse fieldtypes
+# TODO: Implement as `fieldtype(storagetype, :data)`.
+datatype(::Type{<:BlockSparse{<:Any,DataT}}) where {DataT} = DataT
+# TODO: Implement as `ndims(blockoffsetstype(storagetype))`.
+ndims(storagetype::Type{<:BlockSparse{<:Any,<:Any,N}}) where {N} = N
+# TODO: Implement as `fieldtype(storagetype, :blockoffsets)`.
+blockoffsetstype(storagetype::Type{<:BlockSparse}) = BlockOffsets{ndims(storagetype)}
+
+function set_datatype(storagetype::Type{<:BlockSparse}, datatype::Type{<:AbstractVector})
+  return BlockSparse{eltype(datatype),datatype,ndims(storagetype)}
+end
+
+function set_ndims(storagetype::Type{<:BlockSparse}, ndims)
+  return BlockSparse{eltype(storagetype),datatype(storagetype),ndims}
 end
+## end blocksparse fieldtypes
 
 copy(D::BlockSparse) = BlockSparse(copy(data(D)), copy(blockoffsets(D)))
 

NDTensors/src/blocksparse/blocksparsetensor.jl

@@ -42,7 +42,7 @@ function BlockSparseTensor(
 end
 
 function BlockSparseTensor(eltype::Type{<:Number}, blockoffsets::BlockOffsets, inds)
-  return BlockSparseTensor(Vector{eltype}, blockoffsets, inds)
+  return BlockSparseTensor(default_datatype(eltype), blockoffsets, inds)
 end
 
 function BlockSparseTensor(blockoffsets::BlockOffsets, inds)
@@ -61,7 +61,7 @@ function BlockSparseTensor(datatype::Type{<:AbstractArray}, inds)
 end
 
 function BlockSparseTensor(eltype::Type{<:Number}, inds)
-  return BlockSparseTensor(Vector{eltype}, inds)
+  return BlockSparseTensor(default_datatype(eltype), inds)
 end
 
 """
@@ -79,6 +79,14 @@ end
 Construct a block sparse tensor with the specified blocks.
 Defaults to setting structurally non-zero blocks to zero.
 """
+function BlockSparseTensor(
+  datat::Type{<:AbstractArray}, blocks::Vector{BlockT}, inds
+) where {BlockT<:Union{Block,NTuple}}
+  boffs, nnz = blockoffsets(blocks, inds)
+  storage = BlockSparse(datat, boffs, nnz)
+  return tensor(storage, inds)
+end
+
 function BlockSparseTensor(blocks::Vector{BlockT}, inds) where {BlockT<:Union{Block,NTuple}}
   return BlockSparseTensor(Float64, blocks, inds)
 end
@@ -202,6 +210,9 @@ end
 @propagate_inbounds function getindex(
   T::BlockSparseTensor{ElT,N}, i::Vararg{Int,N}
 ) where {ElT,N}
+  if is_unallocated_zeros(T)
+    return zero(ElT)
+  end
   offset, _ = indexoffset(T, i...)
   isnothing(offset) && return zero(ElT)
   return storage(T)[offset]
@@ -225,6 +236,10 @@ end
 function insertblock_offset!(T::BlockSparseTensor{ElT,N}, newblock::Block{N}) where {ElT,N}
   newdim = blockdim(T, newblock)
   newoffset = nnz(T)
+  if newblock ∈ eachnzblock(T)
+    return offset(T, newblock)
+  end
+
   insert!(blockoffsets(T), newblock, newoffset)
   # Insert new block into data
   # TODO: Make GPU-friendly
@@ -849,6 +864,26 @@ function permute_combine(
   return R
 end
 
+function dropzeros(T::Tensor; tol=0)
+  v = Base.similar(NDTensors.data(T), 0)
+  bloc = Vector{Block{ndims(T)}}(undef, 0)
+  off = Vector{Int64}(undef, 0)
+  length = Int64(0)
+  for b in eachnzblock(T)
+    Tb = NDTensors.data(T[b])
+    if norm(Tb) > tol
+      for i in Tb
+        push!(v, i)
+      end
+      #push!(offsets, BlockOffset(b, length))
+      push!(bloc, b)
+      push!(off, length)
+      length += dim(size(Tb))
+    end
+  end
+  return Tensor(BlockSparse(v, Dictionary(bloc, off)), inds(T))
+end
+
 #
 # Print block sparse tensors
 #

NDTensors/src/combiner/contract.jl

@@ -19,6 +19,9 @@ function contract!!(
   tensor::Tensor,
   tensor_labels,
 )
+  output_tensor = allocate(
+    specify_eltype(typeof(output_tensor), eltype(tensor)), output_tensor
+  )
   if ndims(combiner_tensor) ≤ 1
     # Empty combiner, acts as multiplying by 1
     output_tensor = permutedims!!(

NDTensors/src/dense/dense.jl

@@ -147,7 +147,7 @@ function HDF5.write(
   attributes(g)["type"] = "Dense{$(eltype(Store))}"
   attributes(g)["version"] = 1
   if eltype(D) != Nothing
-    write(g, "data", D.data)
+    write(g, "data", NDTensors.allocate(D).data)
   end
 end
 

NDTensors/src/dense/densetensor.jl

@@ -183,6 +183,10 @@ end
 # TODO: Remove this in favor of `map!`
 # applied to `PermutedDimsArray`.
 function permutedims!!(R::DenseTensor, T::DenseTensor, perm, f::Function=(r, t) -> t)
+  ## TODO if unallocatedzero, permute is trivial, just change inds
+  if is_unallocated_zeros(T)
+    return R
+  end
   Base.checkdims_perm(R, T, perm)
   RR = convert(promote_type(typeof(R), typeof(T)), R)
   permutedims!(RR, T, perm, f)

NDTensors/src/dense/fill.jl

@@ -31,6 +31,10 @@ function generic_zeros(
 end
 
 function generic_zeros(StoreT::Type{<:Dense{ElT}}, dim::Integer=0) where {ElT}
+  ## TODO see if this works
+  #  DataT = default_storagetype(ElT)
+  #N = ndims(DataT)
+  #return generic_zeros(UnallocatedZeros{ElT,N,DataT}, dim)
   return generic_zeros(default_storagetype(ElT), dim)
 end