Merge branch 'main' into ranged_truncate_mps

Author: mtfishman

NDTensors/src/diag/diag.jl

@@ -452,7 +452,7 @@ function contract!(
   if all(i -> i < 0, Blabels)
     # If all of B is contracted
     # TODO: can also check NC+NB==NA
-    min_dim = minimum(dims(B))
+    min_dim = min(minimum(dims(A)), minimum(dims(B)))
     if length(Clabels) == 0
       # all indices are summed over, just add the product of the diagonal
       # elements of A and B

src/decomp.jl

@@ -91,7 +91,7 @@ function svd(A::ITensor, Linds...; kwargs...)
   #  @warn "Keyword arguments `utags` and `vtags` are deprecated in favor of `leftags` and `righttags`."
   #end
 
-  Lis = commoninds(A, indices(Linds))
+  Lis = commoninds(A, indices(Linds...))
   Ris = uniqueinds(A, Lis)
 
   Lis_original = Lis
@@ -333,7 +333,7 @@ qr(A::ITensor; kwargs...) = error(noinds_error_message("qr"))
 # call qr on the order-2 tensors directly
 function qr(A::ITensor, Linds...; kwargs...)
   tags::TagSet = get(kwargs, :tags, "Link,qr")
-  Lis = commoninds(A, indices(Linds))
+  Lis = commoninds(A, indices(Linds...))
   Ris = uniqueinds(A, Lis)
 
   Lis_original = Lis
@@ -388,7 +388,7 @@ function factorize_qr(A::ITensor, Linds...; kwargs...)
   if ortho == "left"
     L, R, q = qr(A, Linds...; kwargs...)
   elseif ortho == "right"
-    Lis = uniqueinds(A, indices(Linds))
+    Lis = uniqueinds(A, indices(Linds...))
     R, L, q = qr(A, Lis...; kwargs...)
   else
     error(
@@ -427,9 +427,9 @@ function factorize_eigen(A::ITensor, Linds...; kwargs...)
   ortho::String = get(kwargs, :ortho, "left")
   delta_A2 = get(kwargs, :eigen_perturbation, nothing)
   if ortho == "left"
-    Lis = commoninds(A, indices(Linds))
+    Lis = commoninds(A, indices(Linds...))
   elseif ortho == "right"
-    Lis = uniqueinds(A, indices(Linds))
+    Lis = uniqueinds(A, indices(Linds...))
   else
     error(
       "In factorize using eigen decomposition, ortho keyword $ortho not supported. Supported options are left or right.",
@@ -506,7 +506,8 @@ Note that the default is now `left`, meaning for the results L,R = factorize(A),
   # Determines when to use eigen vs. svd (eigen is less precise,
   # so eigen should only be used if a larger cutoff is requested)
   automatic_cutoff = 1e-12
-  dL, dR = dim(indices(Linds)), dim(indices(setdiff(inds(A), Linds)))
+  Lis = indices(Linds...)
+  dL, dR = dim(Lis), dim(indices(setdiff(inds(A), Lis)))
   maxdim = get(kwargs, :maxdim, min(dL, dR))
   might_truncate = !isnothing(cutoff) || maxdim < min(dL, dR)
 

src/indexset.jl

@@ -23,58 +23,43 @@ const IndexTuple{IndexT<:Index} = Tuple{Vararg{IndexT}}
 # Definition to help with generic code
 const Indices{IndexT<:Index} = Union{Vector{IndexT},Tuple{Vararg{IndexT}}}
 
-# Flatten combinations of tuples and vectors into a single collection
-# of indices
-tuple_vcat(t::Tuple) = t
-tuple_vcat() = ()
-tuple_vcat(t) = (t,)
-tuple_vcat(a, args...) = (tuple_vcat(a)..., tuple_vcat(args...)...)
-
-tuple_to_vector(t::Tuple) = collect(t)
-tuple_to_vector(t) = t
-
-function _narrow_eltype(v::Vector{T}) where {T}
+function _narrow_eltype(v::Vector{T}; default_empty_eltype=T) where {T}
   if isempty(v)
-    return v
+    return default_empty_eltype[]
   end
   return convert(Vector{mapreduce(typeof, promote_type, v)}, v)
 end
-narrow_eltype(v::Vector{T}) where {T} = isconcretetype(T) ? v : _narrow_eltype(v)
-
-push_or_append!(v, x::Union{Vector,Tuple}) = append!(v, x)
-push_or_append!(v, x) = push!(v, x)
-
-function _indices(is::Vector)
-  isempty(is) && return is
-  is_flat = Index[]
-  for i in is
-    push_or_append!(is_flat, i)
-  end
-  return narrow_eltype(is_flat)
-end
-indices(is::Vector{<:Index}) = narrow_eltype(is)
-
-_indices(is::Tuple{Vararg{<:Index}}) = is
-_indices(is::Tuple{Vararg{Union{<:Vector,<:Index}}}) = vcat(is...)
-_indices(is::Tuple{Vararg{Union{<:Tuple,<:Index}}}) = tuple_vcat(is...)
-_indices(is::Tuple{Vararg{Union{<:Tuple,<:Vector,<:Index}}}) = indices(tuple_to_vector.(is))
-indices(is::Tuple{Vararg{<:Index}}) = is
-function indices(is::Tuple)
-  inds = _indices(is)
-  if isempty(inds)
-    # Otherwise it outputs `Any[]`, which breaks
-    # some generic code like `dim`.
-    return Index[]
+function narrow_eltype(v::Vector{T}; default_empty_eltype=T) where {T}
+  if isconcretetype(T)
+    return v
   end
-  return inds
-end
-indices(is::Union{<:Tuple,<:Vector,<:Index}...) = indices(is)
-function indices(is::Vector)
-  # This narrows the type. Also handles the empty case.
-  all(i -> i isa Index, is) && return narrow_eltype(is)
-  return _indices(is)
+  return _narrow_eltype(v; default_empty_eltype)
 end
 
+_indices() = ()
+_indices(x::Index) = (x,)
+
+# Tuples
+_indices(x1::Tuple, x2::Tuple) = (x1..., x2...)
+_indices(x1::Index, x2::Tuple) = (x1, x2...)
+_indices(x1::Tuple, x2::Index) = (x1..., x2)
+_indices(x1::Index, x2::Index) = (x1, x2)
+
+# Vectors
+_indices(x1::Vector, x2::Vector) = narrow_eltype(vcat(x1, x2); default_empty_eltype=Index)
+
+# Mix vectors and tuples/elements
+_indices(x1::Vector, x2) = _indices(x1, [x2])
+_indices(x1, x2::Vector) = _indices([x1], x2)
+_indices(x1::Vector, x2::Tuple) = _indices(x1, [x2...])
+_indices(x1::Tuple, x2::Vector) = _indices([x1...], x2)
+
+indices(x::Vector{Index{S}}) where {S} = x
+indices(x::Vector{Index}) = narrow_eltype(x; default_empty_eltype=Index)
+indices(x::Tuple) = reduce(_indices, x; init=())
+indices(x::Vector) = reduce(_indices, x; init=Index[])
+indices(x...) = indices(x)
+
 # To help with backwards compatibility
 IndexSet(inds::IndexSet) = inds
 IndexSet(inds::Indices) = collect(inds)

src/itensor.jl

@@ -1764,9 +1764,9 @@ T[1, 1, 1] == pT_alias[1, 1, 1]
 ```
 """
 function permute(T::ITensor, new_inds...; kwargs...)
-  if !hassameinds(T, indices(new_inds))
+  if !hassameinds(T, indices(new_inds...))
     error(
-      "In `permute(::ITensor, inds...)`, the input ITensor has indices: \n\n$(inds(T))\n\nbut the desired Index ordering is: \n\n$(indices(new_inds))",
+      "In `permute(::ITensor, inds...)`, the input ITensor has indices: \n\n$(inds(T))\n\nbut the desired Index ordering is: \n\n$(indices(new_inds...))",
     )
   end
   allow_alias = deprecated_keyword_argument(

test/ITensorChainRules/test_chainrules.jl

@@ -352,6 +352,19 @@ Random.seed!(1234)
     @test f2'(x) ≈ 2MPO(s, "I")
     @test f3'(x) ≈ -MPO(s, "I")
   end
+
+  @testset "issue 969" begin
+    i = Index(2)
+    j = Index(3)
+    A = randomITensor(i)
+    B = randomITensor(j)
+    f = function (x, y)
+      d = δ(ind(x, 1), ind(y, 1))
+      return (x * d * y)[]
+    end
+    args = (A, B)
+    test_rrule(ZygoteRuleConfig(), f, args...; rrule_f=rrule_via_ad, check_inferred=false)
+  end
 end
 
 @testset "ChainRules rrules: op" begin

test/ITensorChainRules/test_chainrules_ops.jl

@@ -204,7 +204,9 @@ using Zygote: ZygoteRuleConfig, gradient
     return os
   end
 
-  if VERSION ≥ v"1.7"
+  if VERSION.minor == 7
+    # For some reason this is broken in Julia 1.6 and 1.8?
+    # Seems like a Zygote problem
     f = function (x)
       return ITensor(exp(1.5 * H(x, x); alg=Trotter{1}(1)), s)[1, 1]
     end

test/diagitensor.jl

@@ -530,6 +530,18 @@ using Test
       @test D1 * D2 ≈ dense(D1) * dense(D2)
       @test D2 * D1 ≈ dense(D1) * dense(D2)
     end
+
+    @testset "Rectangular Diag * Dense regression test (#969)" begin
+      i = Index(3)
+      j = Index(2)
+      A = randomITensor(i)
+      B = delta(i, j)
+      C = A * B
+      @test hassameinds(C, j)
+      for n in 1:dim(j)
+        @test C[n] == A[n]
+      end
+    end
   end
 end