Return dimnames as a Vector and the named axes as Tuples, dropping LittleSet (#171)

## Summary

`dimnames` now returns the stored dimension names as a `Vector`, and the
named axes (`inds`, `axes`, `size`) return `Tuple`s, replacing the
`LittleSet` wrapper the accessors previously returned.

`LittleSet` was an ordered-set type whose real jobs were a set-valued
`==` and a custom broadcast style for name-aligned shape combination.
With the accessors returning plain `Vector`s and `Tuple`s, the spots
that relied on the set-valued equality (`==`, `isequal`, `isapprox`, and
the multi-argument `eachindex`) now call `issetequal` directly, and the
broadcast shape machinery (`promote_shape`, `broadcast_shape`,
`check_broadcast_shape`) dispatches on the tuple of named ranges.
`LittleSet` is removed.

Returning `dimnames` as the stored `Vector` exposed a latent autodiff
issue. The accessor carried a `@zero_derivative` Mooncake rule that was
only correct because the old `dimnames` built a fresh object, and
`@zero_derivative` is documented to be unsafe when a function's output
aliases one of its fields. The rule is dropped, so Mooncake
differentiates the field access through its built-in `getfield` rule,
which preserves the aliasing and yields a zero gradient for the
non-differentiable names.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: mtfishman

ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl

@@ -11,7 +11,10 @@ Mooncake.tangent_type(::Type{<:AbstractNamedUnitRange}) = Mooncake.NoTangent
 @zero_derivative DefaultCtx Tuple{typeof(blockedperm), AbstractITensor, Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(blockedperm_nameddims), Any, Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(combine_nameddimsconstructors), Any, Any}
-@zero_derivative DefaultCtx Tuple{typeof(dimnames), Any}
+# `dimnames(::ITensor)` returns the stored names `Vector` directly, so its output
+# aliases a field, where `@zero_derivative` is documented to be incorrect. Let
+# Mooncake differentiate it through the underlying `getfield`, whose built-in rule
+# preserves the aliasing (the names are non-differentiable, so the result is zero).
 @zero_derivative DefaultCtx Tuple{typeof(dimnames), Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(dimnames_setdiff), Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(inds), Any}

src/ITensorBase.jl

@@ -7,7 +7,6 @@ using Compat: @compat
 @compat public @names
 
 # Named-array machinery (relocated from NamedDimsArrays.jl).
-include("littleset.jl")
 include("isnamed.jl")
 include("randname.jl")
 include("abstractnamedinteger.jl")

src/abstractitensor.jl

@@ -53,8 +53,9 @@ denamed(a::AbstractITensor) = throw(MethodError(denamed, a))
 denamed(a::AbstractITensor, inds) = denamed(aligneddims(a, inds))
 dename(a::AbstractITensor, inds) = denamed(aligndims(a, inds))
 
-# Output the named axes/indices of the named dims array.
-inds(a::AbstractITensor) = LittleSet(named.(axes(denamed(a)), dimnames(a)))
+# Output the named axes/indices of the named dims array, as a `Tuple` (even though
+# the dimension names are stored as a `Vector`).
+inds(a::AbstractITensor) = named.(axes(denamed(a)), Tuple(dimnames(a)))
 inds(a::AbstractITensor, dim::Int) = inds(a)[dim]
 
 isnamed(::Type{<:AbstractITensor}) = true
@@ -297,14 +298,6 @@ function Base.similar(
     return similar_nameddims(a, elt, inds)
 end
 
-function Base.similar(a::AbstractArray, inds::LittleSet)
-    return similar_nameddims(a, eltype(a), inds)
-end
-
-function Base.similar(a::AbstractArray, elt::Type, inds::LittleSet)
-    return similar_nameddims(a, elt, inds)
-end
-
 # Same entry points with a named-tensor prototype. An `AbstractITensor` is no longer
 # an `AbstractArray`, so the methods above (which build a named tensor from a plain
 # array prototype) no longer cover it.
@@ -320,11 +313,6 @@ function Base.similar(
     )
     return similar_nameddims(a, elt, inds)
 end
-Base.similar(a::AbstractITensor, inds::LittleSet) = similar_nameddims(a, eltype(a), inds)
-function Base.similar(a::AbstractITensor, elt::Type, inds::LittleSet)
-    return similar_nameddims(a, elt, inds)
-end
-
 function setinds(a::AbstractITensor, inds)
     return nameddimsconstructorof(a)(denamed(a), inds)
 end
@@ -412,10 +400,6 @@ struct NamedDimsCartesianIndices{
         )
     end
 end
-function NamedDimsCartesianIndices(indices::LittleSet)
-    return NamedDimsCartesianIndices(Tuple(indices))
-end
-
 # The element type is no longer carried by the (rank-erased) supertype, so recover
 # it from the stored index-tuple parameter.
 function Base.eltype(
@@ -462,21 +446,21 @@ end
 # Base version ignores dimension names.
 # TODO: Use `mapreduce(isequal, &&, a1, a2)`?
 function Base.isequal(a1::AbstractITensor, a2::AbstractITensor)
-    (inds(a1) == inds(a2)) || return false
+    issetequal(inds(a1), inds(a2)) || return false
     return isequal(denamed(a1), denamed(a2, inds(a1)))
 end
 
 # Base version ignores dimension names.
 # TODO: Use `mapreduce(==, &&, a1, a2)`?
 # TODO: Handle `missing` values properly.
 function Base.:(==)(a1::AbstractITensor, a2::AbstractITensor)
-    (inds(a1) == inds(a2)) || return false
+    issetequal(inds(a1), inds(a2)) || return false
     return denamed(a1) == denamed(a2, inds(a1))
 end
 
 # Base version ignores dimension names.
 function Base.isapprox(a1::AbstractITensor, a2::AbstractITensor; kwargs...)
-    (inds(a1) == inds(a2)) || return false
+    issetequal(inds(a1), inds(a2)) || return false
     return isapprox(denamed(a1), denamed(a2, inds(a1)); kwargs...)
 end
 

src/broadcast.jl

@@ -1,5 +1,5 @@
-using ..ITensorBase: AbstractITensor, AbstractNamedUnitRange, ITensorBase, LittleSet,
-    dename, denamed, getperm, inds, name, named, nameddimsconstructorof
+using ..ITensorBase: AbstractITensor, AbstractNamedUnitRange, ITensorBase, dename, denamed,
+    getperm, inds, name, named, nameddimsconstructorof
 using Base.Broadcast: Broadcast as BC, Broadcasted, broadcast_shape, broadcasted,
     check_broadcast_shape, combine_axes
 using TensorAlgebra: TensorAlgebra as TA
@@ -33,26 +33,41 @@ function BC.combine_axes(a1::AbstractITensor, a2::AbstractITensor)
 end
 BC.combine_axes(a::AbstractITensor) = axes(a)
 
+# The named axes are a `Tuple` of `AbstractNamedUnitRange`s. Dispatch the
+# name-aware shape combination on that tuple form (the elements are not
+# `AbstractUnitRange`s, so Base's positional tuple-shape methods do not apply).
 function BC.broadcast_shape(
-        ax1::LittleSet, ax2::LittleSet, ax_rest::LittleSet...
+        ax1::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}},
+        ax2::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}},
+        ax_rest::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}...
     )
     return broadcast_shape(broadcast_shape(ax1, ax2), ax_rest...)
 end
 
-function BC.broadcast_shape(ax1::LittleSet, ax2::LittleSet)
+function BC.broadcast_shape(
+        ax1::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}},
+        ax2::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}
+    )
     return promote_shape(ax1, ax2)
 end
 
 # Handle scalar values.
-function BC.broadcast_shape(ax1::Tuple{}, ax2::LittleSet)
+function BC.broadcast_shape(
+        ax1::Tuple{}, ax2::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}
+    )
     return ax2
 end
-function BC.broadcast_shape(ax1::LittleSet, ax2::Tuple{})
+function BC.broadcast_shape(
+        ax1::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}, ax2::Tuple{}
+    )
     return ax1
 end
 
-function Base.promote_shape(ax1::LittleSet, ax2::LittleSet)
-    return LittleSet(set_promote_shape(Tuple(ax1), Tuple(ax2)))
+function Base.promote_shape(
+        ax1::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}},
+        ax2::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}
+    )
+    return set_promote_shape(ax1, ax2)
 end
 
 function set_promote_shape(
@@ -83,8 +98,11 @@ function set_promote_shape(
     return ax1
 end
 
-function BC.check_broadcast_shape(ax1::LittleSet, ax2::LittleSet)
-    return set_check_broadcast_shape(Tuple(ax1), Tuple(ax2))
+function BC.check_broadcast_shape(
+        ax1::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}},
+        ax2::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}
+    )
+    return set_check_broadcast_shape(ax1, ax2)
 end
 
 function set_check_broadcast_shape(

src/itensor.jl

@@ -16,9 +16,8 @@ end
 ITensor(a::AbstractITensor, inds) = throw(ArgumentError("Already named."))
 ITensor(a::AbstractITensor) = ITensor(denamed(a), dimnames(a))
 
-# Minimal interface. The dimnames are stored as a `Vector{DimName}`, but the
-# accessor still returns a `LittleSet` over a `Tuple` (unchanged public behavior).
-dimnames(a::ITensor) = LittleSet(Tuple(a.dimnames))
+# Minimal interface. The dimnames are stored as (and returned as) a `Vector`.
+dimnames(a::ITensor) = a.dimnames
 denamed(a::ITensor) = a.denamed
 Base.parent(a::ITensor) = denamed(a)
 

src/littleset.jl

@@ -1,5 +1,5 @@
 using Combinatorics: Combinatorics
-using ITensorBase: @names, AbstractITensor, ITensor, LittleSet, Name, NameMismatch,
+using ITensorBase: @names, AbstractITensor, ITensor, Name, NameMismatch,
     NamedDimsCartesianIndex, NamedDimsCartesianIndices, aligndims, aligneddims, apply,
     dename, denamed, denamedtype, dim, dimnames, dimnametype, dims, fusednames, inds,
     isnamed, mapinds, name, named, nameddims, namedoneto, product, replacedimnames,
@@ -37,7 +37,7 @@ end
         @test inds(na) == (i, j)
         @test inds(na, 1) == i
         @test inds(na, 2) == j
-        @test dimnames(na) == ("i", "j")
+        @test dimnames(na) == ["i", "j"]
         @test dimnames(na, 1) == "i"
         @test dimnames(na, 2) == "j"
         @test dim(na, "i") == 1
@@ -108,9 +108,7 @@ end
         j = namedoneto(4, "j")
         for na′ in (
                 similar(na, Float32, (j, i)),
-                similar(na, Float32, LittleSet((j, i))),
                 similar(a, Float32, (j, i)),
-                similar(a, Float32, LittleSet((j, i))),
             )
             @test eltype(na′) ≡ Float32
             @test all(inds(na′) .== (j, i))
@@ -123,9 +121,7 @@ end
         j = namedoneto(4, "j")
         for na′ in (
                 similar(na, (j, i)),
-                similar(na, LittleSet((j, i))),
                 similar(a, (j, i)),
-                similar(a, LittleSet((j, i))),
             )
             @test eltype(na′) ≡ eltype(na)
             @test all(inds(na′) .== (j, i))
@@ -154,7 +150,7 @@ end
         a = randn(elt, 2)
         na = a[i]
         @test na isa ITensor{String}
-        @test dimnames(na) == ("i",)
+        @test dimnames(na) == ["i"]
         @test denamed(na) == a
 
         # slicing
@@ -416,39 +412,6 @@ end
             @test !iszero(na)
         end
     end
-    @testset "LittleSet" begin
-        # Broadcasting
-        s = LittleSet((1, 2))
-        @test eltype(s) == Int
-        @test s .+ [3, 4] == [4, 6]
-        @test s .+ (3, 4) ≡ (4, 6)
-
-        s = LittleSet(("a", "b", "c"))
-        @test all(s .== ("a", "b", "c"))
-        @test values(s) == ("a", "b", "c")
-        @test Tuple(s) == ("a", "b", "c")
-        @test s[1] == "a"
-        @test s[2] == "b"
-        @test s[3] == "c"
-        for s′ in (
-                replace(x -> x == "b" ? "x" : x, s),
-                replace(s, "b" => "x"),
-                map(x -> x == "b" ? "x" : x, s),
-            )
-            @test s′ isa LittleSet
-            @test Tuple(s′) == ("a", "x", "c")
-            @test s′[1] == "a"
-            @test s′[2] == "x"
-            @test s′[3] == "c"
-        end
-
-        s = LittleSet((1, 2, 3))
-        @test LittleSet(s).values isa Tuple
-        @test LittleSet(s) == s
-        sp = LittleSet{NTuple{3, Float64}}(s)
-        @test eltype(sp) === Float64
-        @test values(sp) == (1.0, 2.0, 3.0)
-    end
     @testset "show" begin
         a = ITensor([1 2; 3 4], ("i", "j"))
         @test sprint(show, "text/plain", a) ==