Keep the trailing auxiliary axis in a dense project

Project.toml

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

src/projectto.jl

@@ -31,16 +31,27 @@ end
     allocate_project(raw, codomain_axes, domain_axes) -> dest
 
 Allocate the destination that projecting `raw` onto
-`codomain_axes`/`domain_axes` fills. The generic method is
-`similar_map(raw, codomain_axes, domain_axes)`. This is a backend
-customization point (with [`projectto!`](@ref) and [`is_projected`](@ref)):
-the allocation may depend on the data, since a symmetric backend derives the
-space of one trailing surplus axis in `raw` (an auxiliary leg appended as
-the last domain axis, e.g. a flux-canceling leg for a charge-shifting
-operator) before allocating.
+`codomain_axes`/`domain_axes` fills. This is a backend customization point
+(with [`projectto!`](@ref) and [`is_projected`](@ref)): the allocation may
+depend on the data, since a trailing surplus axis in `raw` (an auxiliary leg
+appended as the last domain axis, e.g. a flux-canceling leg for a
+charge-shifting operator) has its space taken from `raw` itself on a dense
+backend and derived from the sector structure on a symmetric one.
+
+The generic method keeps that trailing surplus axis (its `raw` axis appended
+to the domain), so the result's rank matches `raw`'s; with no surplus it is
+plain `similar_map(raw, codomain_axes, domain_axes)`.
 """
 function allocate_project(raw, codomain_axes, domain_axes)
-    return similar_map(raw, codomain_axes, domain_axes)
+    nphys = length(codomain_axes) + length(domain_axes)
+    ndims(raw) <= nphys && return similar_map(raw, codomain_axes, domain_axes)
+    ndims(raw) == nphys + 1 || throw(
+        ArgumentError(
+            "`project`: expected at most one trailing auxiliary axis beyond the $nphys \
+            given axes, got a rank-$(ndims(raw)) input"
+        )
+    )
+    return similar_map(raw, codomain_axes, (domain_axes..., axes(raw, nphys + 1)))
 end
 
 """
@@ -130,13 +141,13 @@ tolerances are subject to change in future versions). See
 for the unchecked projection this derives from.
 
 When `raw` has one axis more than the given axes account for, that trailing
-surplus axis is an auxiliary leg whose space a symmetric backend derives so
-the result is symmetry-allowed (e.g. a flux-canceling leg for a
-charge-shifting operator); the result's shape matches `raw`'s shape. The
-derivation is backend-internal: a graded backend reads the sector, the
-`TensorMap` backend projects over the `codomain ⊗ conj(domain)` content. The
-two-argument form takes a flat list of `axes` and is equivalent to an empty
-domain.
+surplus axis is an auxiliary leg appended as the last domain axis, so the
+result's shape matches `raw`'s (e.g. a flux-canceling leg for a
+charge-shifting operator). Its space comes from `raw` itself on a dense
+backend and is derived from the sector structure on a symmetric one (a graded
+backend reads the sector, the `TensorMap` backend projects over the
+`codomain ⊗ conj(domain)` content). The two-argument form takes a flat list
+of `axes` and is equivalent to an empty domain.
 """
 function project(raw, codomain_axes, domain_axes; kwargs...)
     dest = unchecked_project(raw, codomain_axes, domain_axes)

test/test_projectto.jl

@@ -109,3 +109,26 @@ end
     @test size(Msplit) == (2, 3, 1)
     @test vec(Msplit) == vec(flat)
 end
+
+@testset "project keeps a trailing surplus axis ($T)" for T in elts
+    # The dual of the padding case: when `raw` carries one axis *more* than the given axes
+    # account for, that trailing surplus axis is an auxiliary leg (e.g. a flux-canceling leg a
+    # codomain/domain split introduces on a symmetric state/operator), and its space is taken
+    # from `raw`. The result keeps the axis rather than reshaping it away, so its rank matches
+    # `raw`'s, matching the symmetric backends. Here the aux leg is the dim-1 leg a caller adds
+    # with `reshape(a, (size(a)..., 1))`.
+    raw = randn(T, 2, 3, 1)
+
+    # all-codomain (state) form
+    M = project(raw, (Base.OneTo(2), Base.OneTo(3)))
+    @test size(M) == (2, 3, 1)
+    @test M == raw
+
+    # explicit split: the surplus axis lands past the codomain/domain axes given
+    Msplit = project(raw, (Base.OneTo(2),), (Base.OneTo(3),))
+    @test size(Msplit) == (2, 3, 1)
+    @test Msplit == raw
+
+    # more than one surplus axis is rejected
+    @test_throws ArgumentError project(randn(T, 2, 3, 1, 1), (Base.OneTo(2), Base.OneTo(3)))
+end