Loosen generic tensor operations to accept non-array operands

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.14.1"
+version = "0.15.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

docs/Project.toml

@@ -11,4 +11,4 @@ path = ".."
 Documenter = "1.8.1"
 ITensorFormatter = "0.2.27"
 Literate = "2.20.1"
-TensorAlgebra = "0.14"
+TensorAlgebra = "0.15"

examples/Project.toml

@@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 path = ".."
 
 [compat]
-TensorAlgebra = "0.14"
+TensorAlgebra = "0.15"

ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl

@@ -54,9 +54,9 @@ end
 # in-place
 function TA.contractopadd!(
         algorithm::TensorOperationsAlgorithm,
-        a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain,
-        op1, a1::AbstractArray, perm1_codomain, perm1_domain,
-        op2, a2::AbstractArray, perm2_codomain, perm2_domain,
+        a_dest, perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain,
         α::Number, β::Number
     )
     permblocks1 = Tuple.((perm1_codomain, perm1_domain))

src/TensorAlgebra.jl

@@ -8,7 +8,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public biperm, bipartition, contractopadd!, label_type, matricizeop, to_range, zero!, scale!, permuteddims, conjed, ConjArray"
+            "public biperm, bipartition, contractopadd!, label_type, matricizeopperm, to_range, zero!, scale!, permuteddims, conjed, ConjArray"
         )
     )
 end

src/contract/allocate_output.jl

@@ -1,5 +1,3 @@
-using Base.PermutedDimsArrays: genperm
-
 function check_biperm(a, perm_codomain, perm_domain)
     ndims(a) == length(perm_codomain) + length(perm_domain) ||
         throw(ArgumentError("Invalid bipartitioned permutation"))
@@ -44,25 +42,22 @@ end
 function output_axes(
         ::typeof(contract),
         perm_dest_codomain, perm_dest_domain,
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain
-    )
-    axes_codomain, axes_contracted = bipartition(axes(a1), perm1_codomain, perm1_domain)
-    axes_contracted2, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain)
-    @assert length.(axes_contracted) == length.(axes_contracted2)
-    # default: flatten the destination permutation
-    return genperm(
-        (axes_codomain..., axes_domain...), (perm_dest_codomain..., perm_dest_domain...)
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain
     )
+    axes_codomain, _ = bipartition(axes(a1), perm1_codomain, perm1_domain)
+    _, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain)
+    axes_uncontracted = (axes_codomain..., axes_domain...)
+    return bipartition(axes_uncontracted, perm_dest_codomain, perm_dest_domain)
 end
 
 # TODO: Use `ArrayLayouts`-like `MulAdd` object,
 # i.e. `ContractAdd`?
 function allocate_output(
         ::typeof(contract),
         perm_dest_codomain, perm_dest_domain,
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain
     )
     check_input(
         contract,
@@ -73,12 +68,14 @@ function allocate_output(
         perm2_codomain,
         perm2_domain
     )
-    axes_dest = output_axes(
+    codomain_axes_dest, domain_axes_dest = output_axes(
         contract,
         perm_dest_codomain, perm_dest_domain,
         a1, perm1_codomain, perm1_domain,
         a2, perm2_codomain, perm2_domain
     )
     T = promote_type(eltype(a1), eltype(a2))
-    return zero!(similar(a1, T, axes_dest))
+    # `domain_axes_dest` come straight from `axes(a2)` (stored/dualized convention), so
+    # un-dualize them into `similar_map`'s codomain-facing convention.
+    return zero!(similar_map(a1, T, codomain_axes_dest, conj.(domain_axes_dest)))
 end

src/contract/contract.jl

@@ -2,7 +2,7 @@
 # TODO: Add `scaledcontract(a1, labels1, a2, labels2, α) = α * contract(a1, labels1, a2, labels2)`.
 
 # contract (labels)
-function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs...)
+function contract(a1, labels1, a2, labels2; kwargs...)
     # Optionally convert the labels to a representation cheaper to run the bookkeeping on (see
     # `label_type`). `encode_contraction_labels`/`decode_contraction_labels` are no-ops unless the label type opts in.
     l1, l2 = encode_contraction_labels(labels1, labels2)
@@ -11,7 +11,7 @@ function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs
     return a_dest, decode_contraction_labels(l_dest, labels1, labels2)
 end
 function contract(
-        labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs...
+        labels_dest, a1, labels1, a2, labels2; kwargs...
     )
     t1 = ntuple(i -> labels1[i], Val(ndims(a1)))
     t2 = ntuple(i -> labels2[i], Val(ndims(a2)))
@@ -30,7 +30,7 @@ function contract(
     )
 end
 function _contract(
-        ::Val{K}, labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2,
+        ::Val{K}, labels_dest, a1, labels1, a2, labels2,
         contracted1; kwargs...
     ) where {K}
     biperm_dest, biperm1, biperm2 =
@@ -40,8 +40,8 @@ end
 
 # contract (bipartitioned permutations)
 function contract(
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain;
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain;
         kwargs...
     )
     Ndest_codomain = Val(length(perm1_codomain))
@@ -57,8 +57,8 @@ function contract(
 end
 function contract(
         perm_dest_codomain, perm_dest_domain,
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain;
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain;
         kwargs...
     )
     a_dest = allocate_output(
@@ -77,19 +77,19 @@ end
 
 # contract! (labels)
 function contract!(
-        a_dest::AbstractArray, labels_dest,
-        a1::AbstractArray, labels1,
-        a2::AbstractArray, labels2;
+        a_dest, labels_dest,
+        a1, labels1,
+        a2, labels2;
         kwargs...
     )
     return contractadd!(
         a_dest, labels_dest, a1, labels1, a2, labels2, true, false; kwargs...
     )
 end
 function contract!(
-        a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain,
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain;
+        a_dest, perm_dest_codomain, perm_dest_domain,
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain;
         kwargs...
     )
     return contractadd!(
@@ -102,9 +102,9 @@ end
 
 # contractadd! (labels)
 function contractadd!(
-        a_dest::AbstractArray, labels_dest,
-        a1::AbstractArray, labels1,
-        a2::AbstractArray, labels2,
+        a_dest, labels_dest,
+        a1, labels1,
+        a2, labels2,
         α::Number, β::Number;
         kwargs...
     )
@@ -114,9 +114,9 @@ function contractadd!(
 end
 # contractadd! (bipartitioned permutations)
 function contractadd!(
-        a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain,
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain,
+        a_dest, perm_dest_codomain, perm_dest_domain,
+        a1, perm1_codomain, perm1_domain,
+        a2, perm2_codomain, perm2_domain,
         α::Number, β::Number;
         kwargs...
     )
@@ -130,9 +130,9 @@ end
 
 # contractopadd! (labels)
 function contractopadd!(
-        a_dest::AbstractArray, labels_dest,
-        op1, a1::AbstractArray, labels1,
-        op2, a2::AbstractArray, labels2,
+        a_dest, labels_dest,
+        op1, a1, labels1,
+        op2, a2, labels2,
         α::Number, β::Number;
         kwargs...
     )
@@ -147,8 +147,8 @@ function contractopadd!(
     )
 end
 function _contractopadd!(
-        ::Val{K}, a_dest::AbstractArray, labels_dest,
-        op1, a1::AbstractArray, labels1, op2, a2::AbstractArray, labels2,
+        ::Val{K}, a_dest, labels_dest,
+        op1, a1, labels1, op2, a2, labels2,
         α::Number, β::Number, contracted1; kwargs...
     ) where {K}
     biperm_dest, biperm1, biperm2 =
@@ -159,9 +159,9 @@ function _contractopadd!(
 end
 # contractopadd! (bipartitioned permutations, algorithm selection)
 function contractopadd!(
-        a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain,
-        op1, a1::AbstractArray, perm1_codomain, perm1_domain,
-        op2, a2::AbstractArray, perm2_codomain, perm2_domain,
+        a_dest, perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain,
         α::Number, β::Number;
         alg = DefaultContractAlgorithm(), kwargs...
     )
@@ -184,9 +184,9 @@ end
 # Required interface if not using matricized contraction
 function contractopadd!(
         algorithm::ContractAlgorithm,
-        a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain,
-        op1, a1::AbstractArray, perm1_codomain, perm1_domain,
-        op2, a2::AbstractArray, perm2_codomain, perm2_domain,
+        a_dest, perm_dest_codomain, perm_dest_domain,
+        op1, a1, perm1_codomain, perm1_domain,
+        op2, a2, perm2_codomain, perm2_domain,
         α::Number, β::Number
     )
     return throw(

src/contract/contract_matricize.jl

@@ -16,20 +16,22 @@ function contractopadd!(
         a1, biperm1_codomain, biperm1_domain,
         a2, biperm2_codomain, biperm2_domain
     )
-    a1_mat =
-        matricizeop(algorithm.left_fusion_style, op1, a1, biperm1_codomain, biperm1_domain)
-    a2_mat =
-        matricizeop(algorithm.right_fusion_style, op2, a2, biperm2_codomain, biperm2_domain)
+    a1_mat = matricizeopperm(
+        algorithm.left_fusion_style, op1, a1, biperm1_codomain, biperm1_domain
+    )
+    a2_mat = matricizeopperm(
+        algorithm.right_fusion_style, op2, a2, biperm2_codomain, biperm2_domain
+    )
     output_style = algorithm.output_fusion_style
     # Matricize the destination and multiply straight into it: a no-op for an aligned or
     # transposed dense output (a view aliasing `a_dest`, so `mul!` writes through and we
     # are done), a fresh permuted copy otherwise. Either way `matricize` seeds `a_dest_mat`
     # with `a_dest`'s current contents, so `β` rides on the `mul!` and a detached copy is
     # written back with a plain overwrite.
-    a_dest_mat = matricize(output_style, a_dest, invperm_codomain, invperm_domain)
+    a_dest_mat = matricizeperm(output_style, a_dest, invperm_codomain, invperm_domain)
     mul!(a_dest_mat, a1_mat, a2_mat, α, β)
     if !Base.mightalias(a_dest_mat, a_dest)
-        unmatricize!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain)
+        unmatricizeperm!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain)
     end
     return a_dest
 end

src/contract/contractalgorithm.jl

@@ -11,20 +11,16 @@ end
 Matricize(fusion_style) = Matricize(fusion_style, fusion_style, fusion_style)
 Matricize() = Matricize(ReshapeFusion())
 
-function select_contract_algorithm(algorithm, a1::AbstractArray, a2::AbstractArray)
+function select_contract_algorithm(algorithm, a1, a2)
     return error("Not implemented.")
 end
-function select_contract_algorithm(
-        algorithm::ContractAlgorithm, a1::AbstractArray, a2::AbstractArray
-    )
+function select_contract_algorithm(algorithm::ContractAlgorithm, a1, a2)
     return algorithm
 end
-function select_contract_algorithm(
-        algorithm::DefaultContractAlgorithm, a1::AbstractArray, a2::AbstractArray
-    )
+function select_contract_algorithm(algorithm::DefaultContractAlgorithm, a1, a2)
     return default_contract_algorithm(a1, a2)
 end
-function default_contract_algorithm(a1::AbstractArray, a2::AbstractArray)
+function default_contract_algorithm(a1, a2)
     return default_contract_algorithm(typeof(a1), typeof(a2))
 end
 function default_contract_algorithm(A1::Type{<:AbstractArray}, A2::Type{<:AbstractArray})