Small PR to start getting Dagger to work in the Dense framework

Author: kmp5VT

NDTensors/ext/NDTensorsDaggerExt/NDTensorsDaggerExt.jl

@@ -0,0 +1,4 @@
+module NDTensorsDaggerExt
+include("set_types.jl")
+include("similar.jl")
+end

NDTensors/ext/NDTensorsDaggerExt/set_types.jl

@@ -0,0 +1,23 @@
+# TypeParameterAccessors definitions
+using Dagger: Dagger, DArray
+using NDTensors.TypeParameterAccessors:
+  TypeParameterAccessors, Position, default_type_parameters, parameter, position
+
+blocktype(darray::DArray) = blocktype(typeof(darray))
+blocktype(darrayT::Type{<:DArray}) = parameter(darrayT, position(darrayT, blocktype))
+
+function TypeParameterAccessors.position(::Type{<:DArray}, ::typeof(blocktype))
+  return Position(3)
+end
+
+concattype(darray::DArray) = concattype(typeof(darray))
+concattype(darrayT::Type{<:DArray}) = parameter(darrayT, position(darrayT, 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

NDTensors/ext/NDTensorsDaggerExt/similar.jl

@@ -0,0 +1,12 @@
+using NDTensors: NDTensors
+using NDTensors.Expose: Exposed, unexpose
+using Dagger: DArray
+
+function NDTensors.similar(E::Exposed{<:DArray})
+  A = unexpose(E)
+  return Base.similar(A)
+end
+
+function similar(E::Exposed{<:DArray}, eltype::Type)
+  return Base.similar(unexpose(E), eltype)
+end

NDTensors/src/abstractarray/similar.jl

@@ -87,6 +87,9 @@ end
 # NDTensors.similar
 similar(array::AbstractArray, dims::Tuple) = NDTensors.similar(typeof(array), dims)
 
+function similar(E::Exposed, eltype::Type)
+  return similar(unexpose(E), eltype)
+end
 # Use the `size` to determine the dimensions
 # NDTensors.similar
 function similar(array::AbstractArray, eltype::Type)
@@ -95,6 +98,7 @@ end
 
 # Use the `size` to determine the dimensions
 # NDTensors.similar
+similar(E::Exposed) = similar(unexpose(E))
 similar(array::AbstractArray) = NDTensors.similar(typeof(array), size(array))
 
 ## similartype

NDTensors/src/tensorstorage/similar.jl

@@ -1,9 +1,9 @@
 # NDTensors.similar
-similar(storage::TensorStorage) = setdata(storage, NDTensors.similar(data(storage)))
+similar(storage::TensorStorage) = setdata(storage, NDTensors.similar(expose(data(storage))))
 
 # NDTensors.similar
 function similar(storage::TensorStorage, eltype::Type)
-  return setdata(storage, NDTensors.similar(data(storage), eltype))
+  return setdata(storage, NDTensors.similar(expose(data(storage)), eltype))
 end
 
 # NDTensors.similar