Compare BiTuple by its flat contents

Project.toml

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

src/bituple.jl

@@ -29,11 +29,27 @@ function Base.show(io::IO, bt::BiTuple)
     return print(io, "BiTuple(", bt.t1, ", ", bt.t2, ")")
 end
 
-Base.:(==)(a::BiTuple, b::BiTuple) = a.t1 == b.t1 && a.t2 == b.t2
-Base.hash(bt::BiTuple, h::UInt) = hash(bt.t2, hash(bt.t1, hash(:BiTuple, h)))
+# Equality ignores the split: it compares the flattened contents, so a `(2, 0)` and a `(1, 1)` split
+# with the same flat entries are equal, and a `BiTuple` equals the plain tuple it flattens to. This
+# mirrors how comparing arrays ignores how they are partitioned into blocks. `hash` matches, hashing
+# the flat form so the equality relation stays hash-consistent.
+Base.:(==)(a::BiTuple, b::BiTuple) = Tuple(a) == Tuple(b)
+Base.:(==)(bt::BiTuple, t::Tuple) = Tuple(bt) == t
+Base.:(==)(t::Tuple, bt::BiTuple) = t == Tuple(bt)
+Base.hash(bt::BiTuple, h::UInt) = hash(Tuple(bt), h)
 
 Base.invperm(bt::BiTuple{N1}) where {N1} = BiTuple(invperm(Tuple(bt)), Val(N1))
 
+# A single-argument `map` preserves the split, mapping each block: with one operand there is no
+# ambiguity about which split to keep. (The multi-argument case is the ambiguous one and is left to
+# collapse to the flat tuple.)
+Base.map(f, bt::BiTuple) = BiTuple(map(f, bt.t1), map(f, bt.t2))
+bipartition_axes(bt::BiTuple, split...) = bipartition_axes(Tuple(bt), split...)
+bipartition(bt::BiTuple, length1::Val) = bipartition(Tuple(bt), length1)
+function bipartition(bt::BiTuple, group1::Tuple, group2::Tuple)
+    return bipartition(Tuple(bt), group1, group2)
+end
+
 """
     bipartition(t::Tuple, length1::Val) -> (t1, t2)
     bipartition(t::Tuple, group1::Tuple, group2::Tuple) -> (p1, p2)

src/linearbroadcasted.jl

@@ -128,11 +128,13 @@ addends(a::AddBroadcasted) = a.args
 # differing shapes), so combine by verifying equality through `axes` (TensorAlgebra's, which
 # works for a non-`AbstractArray` backend like a `TensorMap`) rather than Base's `combine_axes`,
 # which would call `Base.axes`/`Base.size` on the operands. A mismatch (e.g. a half-conjugated
-# `conj.(a) .- b`, whose dualized and non-dualized axes differ) throws here.
+# `conj.(a) .- b`, whose dualized and non-dualized axes differ) throws here. Axis equality ignores
+# the codomain/domain split, so operands that agree on the flat legs but differ in split still
+# combine, matching the split-collapsing result.
 function Base.axes(a::AddBroadcasted)
     axs = map(axes, addends(a))
     ax = first(axs)
-    all(x -> x == ax, axs) ||
+    all(==(ax), axs) ||
         throw(DimensionMismatch("linear-combination operands have mismatched axes: $axs"))
     return ax
 end

src/matricize.jl

@@ -255,7 +255,7 @@ function unmatricizeperm(
         throw(ArgumentError("axes do not match permutation"))
     codomain_axes, domain_axes = bipartition_axes(axes_dest, invbiperm)
     a12 = unmatricize(style, m, codomain_axes, domain_axes)
-    biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(axes_dest)))
+    biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(invbiperm)))
     return bipermutedims(a12, biperm_dest)
 end
 
@@ -274,7 +274,7 @@ function unmatricizeperm!(
         throw(ArgumentError("destination does not match permutation"))
     codomain_axes, domain_axes = bipartition_axes(axes(a_dest), invbiperm)
     a_perm = unmatricize(style, m, codomain_axes, domain_axes)
-    biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(axes(a_dest))))
+    biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(invbiperm)))
     return bipermutedims!(a_dest, a_perm, biperm_dest)
 end
 

test/test_bituple.jl

@@ -25,9 +25,22 @@ using TestExtras: @constinferred
     # Split constructor: split a flat tuple at the given codomain length.
     @test (@constinferred BiTuple((3, 4, 5, 2, 1), Val(3))) == BiTuple((3, 4, 5), (2, 1))
 
-    # Equality compares the two blocks.
+    # Equality ignores the split, comparing the flattened contents: same flat entries compare equal
+    # even under a different codomain/domain split, the way array equality ignores block partitioning.
     @test BiTuple((1, 2), (3,)) == BiTuple((1, 2), (3,))
-    @test BiTuple((1, 2), (3,)) != BiTuple((1,), (2, 3))
+    @test BiTuple((1, 2), (3,)) == BiTuple((1,), (2, 3))
+    @test BiTuple((1, 2), (3,)) != BiTuple((1, 3), (2,))
+
+    # A `BiTuple` also equals the plain tuple it flattens to.
+    @test BiTuple((1, 2), (3,)) == (1, 2, 3)
+    @test (1, 2, 3) == BiTuple((1, 2), (3,))
+    @test BiTuple((1, 2), (3,)) != (1, 2)
+    # `hash` matches equality (hashes the flat form).
+    @test hash(BiTuple((1, 2), (3,))) == hash(BiTuple((1,), (2, 3))) == hash((1, 2, 3))
+
+    # Single-argument `map` preserves the split (no ambiguity with one operand). Checked with `===`
+    # since `==` ignores the split and so would not catch a collapse to the flat form.
+    @test (@constinferred map(x -> x + 1, BiTuple((1, 2), (3,)))) === BiTuple((2, 3), (4,))
 end
 
 @testset "biperm" begin