Updates to get contract working

Author: kmp5VT

NDTensors/ext/NDTensorsDaggerExt/set_types.jl

@@ -1,7 +1,8 @@
 # TypeParameterAccessors definitions
-using Dagger: Dagger, DArray
+using Dagger: Dagger, Blocks, DArray
+using NDTensors: NDTensors
 using NDTensors.TypeParameterAccessors:
-  TypeParameterAccessors, Position, default_type_parameters, parameter, position
+  TypeParameterAccessors, Position, default_type_parameters, parameter, position, set_type_parameters
 
 blocktype(darray::DArray) = blocktype(typeof(darray))
 blocktype(darrayT::Type{<:DArray}) = parameter(darrayT, position(darrayT, blocktype))
@@ -13,11 +14,16 @@ end
 concattype(darray::DArray) = concattype(typeof(darray))
 concattype(darrayT::Type{<:DArray}) = parameter(darrayT, position(darrayT, concattype))
 
-function TypeParameterAccessors.position(::Type{<:DArray}, ::typeof(concattype))
+function TypeParameterAccessors.position(::Type{<:DArray}, ::typeof(concattype)) 
   return Position(4)
 end
 
 ## TODO use autoblock
 function TypeParameterAccessors.default_type_parameters(::Type{<:DArray})
   return (default_type_parameters(AbstractArray)..., Blocks{2}, typeof(cat))
 end
+
+## TODO need to make this work. Need to specify 
+function NDTensors.set_ndims(type::Type{<:DArray}, param)
+  return set_type_parameters(type, (ndims, blocktype), (param, Blocks{param}))
+end

NDTensors/src/tensorstorage/similar.jl

@@ -45,6 +45,7 @@ function similar(storagetype::Type{<:TensorStorage}, dims::Dims)
   # TODO: Don't convert to an `AbstractVector` with `prod`, once we support
   # more general data types.
   # return setdata(storagetype, NDTensors.similar(datatype(storagetype), dims))
+  ## TODO use expose here to work for DArray
   return setdata(storagetype, NDTensors.similar(datatype(storagetype), prod(dims)))
 end