Format

ITensor · Oct 23, 2023 · d60b852 · d60b852
1 parent c133f40
commit d60b852
Show file tree

Hide file tree

Showing 8 changed files with 35 additions and 9 deletions.
diff --git a/NDTensors/ext/NDTensorsMetalExt/linearalgebra.jl b/NDTensors/ext/NDTensorsMetalExt/linearalgebra.jl
@@ -10,5 +10,7 @@ end
 
 function NDTensors.svd(leaf_parenttype::Type{<:MtlArray}, A::AbstractMatrix)
   U, S, V = NDTensors.svd(NDTensors.cpu(A))
-  return adapt(leaf_parenttype, U), adapt(set_ndims(leaf_parenttype, ndims(S)), S), adapt(leaf_parenttype, V)
+  return adapt(leaf_parenttype, U),
+  adapt(set_ndims(leaf_parenttype, ndims(S)), S),
+  adapt(leaf_parenttype, V)
 end
diff --git a/NDTensors/ext/NDTensorsMetalExt/permutedims.jl b/NDTensors/ext/NDTensorsMetalExt/permutedims.jl
@@ -1,4 +1,10 @@
-function NDTensors.permutedims!(::Type{<:MtlArray}, Adest::Base.ReshapedArray{<:Any,<:Any,<:SubArray}, ::Type{<:MtlArray}, A, perm)
+function NDTensors.permutedims!(
+  ::Type{<:MtlArray},
+  Adest::Base.ReshapedArray{<:Any,<:Any,<:SubArray},
+  ::Type{<:MtlArray},
+  A,
+  perm,
+)
   Aperm = permutedims(A, perm)
   Adest_parent = parent(Adest)
   copyto!(Adest_parent, Aperm)

diff --git a/NDTensors/src/abstractarray/permutedims.jl b/NDTensors/src/abstractarray/permutedims.jl
@@ -26,7 +26,9 @@ function permutedims!!(B::AbstractArray, A::AbstractArray, perm)
   return permutedims!!(leaf_parenttype(B), B, leaf_parenttype(A), A, perm)
 end
 
-function permutedims!!(Bleaftype::Type{<:AbstractArray}, B, Aleaftype::Type{<:AbstractArray}, A, perm)
+function permutedims!!(
+  Bleaftype::Type{<:AbstractArray}, B, Aleaftype::Type{<:AbstractArray}, A, perm
+)
   permutedims!(leaf_parenttype(B), B, leaf_parenttype(A), A, perm)
   return B
 end

diff --git a/NDTensors/src/blocksparse/blocksparse.jl b/NDTensors/src/blocksparse/blocksparse.jl
@@ -55,7 +55,11 @@ function BlockSparse(
 end
 
 function BlockSparse(
-  datatype::Type{<:AbstractArray}, ::UndefInitializer, blockoffsets::BlockOffsets, dim::Integer; vargs...
+  datatype::Type{<:AbstractArray},
+  ::UndefInitializer,
+  blockoffsets::BlockOffsets,
+  dim::Integer;
+  vargs...,
 )
   return BlockSparse(datatype(undef, dim), blockoffsets; vargs...)
 end

diff --git a/NDTensors/src/blocksparse/blocksparsetensor.jl b/NDTensors/src/blocksparse/blocksparsetensor.jl
@@ -32,7 +32,10 @@ function BlockSparseTensor(
 end
 
 function BlockSparseTensor(
-  datatype::Type{<:AbstractArray}, ::UndefInitializer, blocks::Vector{<:Union{Block,NTuple}}, inds
+  datatype::Type{<:AbstractArray},
+  ::UndefInitializer,
+  blocks::Vector{<:Union{Block,NTuple}},
+  inds,
 )
   boffs, nnz = blockoffsets(blocks, inds)
   storage = BlockSparse(datatype, undef, boffs, nnz)

diff --git a/NDTensors/src/blocksparse/diagblocksparse.jl b/NDTensors/src/blocksparse/diagblocksparse.jl
@@ -63,7 +63,10 @@ function DiagBlockSparse(
 end
 
 function DiagBlockSparse(
-  datatype::Type{<:AbstractArray}, ::UndefInitializer, boffs::BlockOffsets, diaglength::Integer
+  datatype::Type{<:AbstractArray},
+  ::UndefInitializer,
+  boffs::BlockOffsets,
+  diaglength::Integer,
 )
   return DiagBlockSparse(datatype(undef, diaglength), boffs)
 end

diff --git a/NDTensors/src/blocksparse/linearalgebra.jl b/NDTensors/src/blocksparse/linearalgebra.jl
@@ -148,7 +148,9 @@ function svd(T::BlockSparseMatrix{ElT}; kwargs...) where {ElT}
   indsS = setindex(indsS, dag(vind), 2)
 
   U = BlockSparseTensor(leaf_parenttype(T), undef, nzblocksU, indsU)
-  S = DiagBlockSparseTensor(set_eltype(leaf_parenttype(T), real(ElT)), undef, nzblocksS, indsS)
+  S = DiagBlockSparseTensor(
+    set_eltype(leaf_parenttype(T), real(ElT)), undef, nzblocksS, indsS
+  )
   V = BlockSparseTensor(leaf_parenttype(T), undef, nzblocksV, indsV)
 
   for n in 1:nnzblocksT
@@ -302,7 +304,9 @@ function eigen(
     nzblocksV[n] = blockV
   end
 
-  D = DiagBlockSparseTensor(set_ndims(set_eltype(leaf_parenttype(T), ElD), 1), undef, nzblocksD, indsD)
+  D = DiagBlockSparseTensor(
+    set_ndims(set_eltype(leaf_parenttype(T), ElD), 1), undef, nzblocksD, indsD
+  )
   V = BlockSparseTensor(set_eltype(leaf_parenttype(T), ElV), undef, nzblocksV, indsV)
 
   for n in 1:nnzblocksT

diff --git a/NDTensors/src/dense/densetensor.jl b/NDTensors/src/dense/densetensor.jl
@@ -119,7 +119,9 @@ end
   return T
 end
 
-@propagate_inbounds function setindex!(::Type{<:AbstractArray}, T::DenseTensor{<:Number}, x::Number)
+@propagate_inbounds function setindex!(
+  ::Type{<:AbstractArray}, T::DenseTensor{<:Number}, x::Number
+)
   data(T)[] = x
   return T
 end