From c49d7f2f9d1c3ceb0fd07b1ab7181b7b06dd5981 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= <ogauthe@flatironinstitute.org>
Date: Tue, 5 Nov 2024 14:26:10 -0500
Subject: [PATCH] [GradedAxes] Introduce GradedUnitRangeDual (#1531)

---
 .../test/runtests.jl                          | 157 ++++++++++++--
 .../src/abstractblocksparsearray/views.jl     |  13 +-
 .../lib/BlockSparseArrays/test/test_basics.jl |  21 +-
 .../src/lib/GradedAxes/src/GradedAxes.jl      |   3 +-
 .../lib/GradedAxes/src/blockedunitrange.jl    |   2 +-
 NDTensors/src/lib/GradedAxes/src/dual.jl      |   7 +-
 NDTensors/src/lib/GradedAxes/src/fusion.jl    |  15 +-
 .../src/lib/GradedAxes/src/gradedunitrange.jl | 176 ++++++++--------
 .../lib/GradedAxes/src/gradedunitrangedual.jl | 105 ++++++++++
 NDTensors/src/lib/GradedAxes/src/onetoone.jl  |   9 +
 .../src/lib/GradedAxes/src/unitrangedual.jl   | 135 -------------
 .../src/lib/GradedAxes/test/test_basics.jl    |  58 +++++-
 .../src/lib/GradedAxes/test/test_dual.jl      | 191 ++++++++++++------
 .../GradedAxes/test/test_tensor_product.jl    |   4 +-
 .../LabelledNumbers/src/labelledunitrange.jl  |  10 -
 15 files changed, 585 insertions(+), 321 deletions(-)
 create mode 100644 NDTensors/src/lib/GradedAxes/src/gradedunitrangedual.jl
 create mode 100644 NDTensors/src/lib/GradedAxes/src/onetoone.jl
 delete mode 100644 NDTensors/src/lib/GradedAxes/src/unitrangedual.jl

diff --git a/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl b/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl
index 38142b65f5..585dbb5afe 100644
--- a/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl
+++ b/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl
@@ -1,13 +1,22 @@
 @eval module $(gensym())
 using Compat: Returns
 using Test: @test, @testset, @test_broken
-using BlockArrays: Block, BlockedOneTo, blockedrange, blocklengths, blocksize
+using BlockArrays:
+  AbstractBlockArray, Block, BlockedOneTo, blockedrange, blocklengths, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
 using NDTensors.GradedAxes:
-  GradedAxes, GradedOneTo, UnitRangeDual, blocklabels, dual, gradedrange
+  GradedAxes,
+  GradedOneTo,
+  GradedUnitRange,
+  GradedUnitRangeDual,
+  blocklabels,
+  dual,
+  gradedrange,
+  isdual
 using NDTensors.LabelledNumbers: label
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: fusedims, splitdims
+using LinearAlgebra: adjoint
 using Random: randn!
 function blockdiagonal!(f, a::AbstractArray)
   for i in 1:minimum(blocksize(a))
@@ -31,6 +40,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     d2 = gradedrange([U1(0) => 2, U1(1) => 2])
     a = BlockSparseArray{elt}(d1, d2, d1, d2)
     blockdiagonal!(randn!, a)
+    @test axes(a, 1) isa GradedOneTo
+    @test axes(view(a, 1:4, 1:4, 1:4, 1:4), 1) isa GradedOneTo
 
     for b in (a + a, 2 * a)
       @test size(b) == (4, 4, 4, 4)
@@ -38,8 +49,6 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test blocklengths.(axes(b)) == ([2, 2], [2, 2], [2, 2], [2, 2])
       @test nstored(b) == 32
       @test block_nstored(b) == 2
-      # TODO: Have to investigate why this fails
-      # on Julia v1.6, or drop support for v1.6.
       for i in 1:ndims(a)
         @test axes(b, i) isa GradedOneTo
       end
@@ -103,6 +112,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test blocksize(m) == (3, 3)
     @test a == splitdims(m, (d1, d2), (d1, d2))
   end
+
   @testset "dual axes" begin
     r = gradedrange([U1(0) => 2, U1(1) => 2])
     for ax in ((r, r), (dual(r), r), (r, dual(r)), (dual(r), dual(r)))
@@ -110,9 +120,9 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @views for b in [Block(1, 1), Block(2, 2)]
         a[b] = randn(elt, size(a[b]))
       end
-      # TODO: Define and use `isdual` here.
       for dim in 1:ndims(a)
         @test typeof(ax[dim]) === typeof(axes(a, dim))
+        @test isdual(ax[dim]) == isdual(axes(a, dim))
       end
       @test @view(a[Block(1, 1)])[1, 1] == a[1, 1]
       @test @view(a[Block(1, 1)])[2, 1] == a[2, 1]
@@ -130,17 +140,16 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
         @test a[I] == a_dense[I]
       end
       @test axes(a') == dual.(reverse(axes(a)))
-      # TODO: Define and use `isdual` here.
-      @test typeof(axes(a', 1)) === typeof(dual(axes(a, 2)))
-      @test typeof(axes(a', 2)) === typeof(dual(axes(a, 1)))
+
+      @test isdual(axes(a', 1)) ≠ isdual(axes(a, 2))
+      @test isdual(axes(a', 2)) ≠ isdual(axes(a, 1))
       @test isnothing(show(devnull, MIME("text/plain"), a))
 
       # Check preserving dual in tensor algebra.
       for b in (a + a, 2 * a, 3 * a - a)
         @test Array(b) ≈ 2 * Array(a)
-        # TODO: Define and use `isdual` here.
         for dim in 1:ndims(a)
-          @test typeof(axes(b, dim)) === typeof(axes(b, dim))
+          @test isdual(axes(b, dim)) == isdual(axes(a, dim))
         end
       end
 
@@ -148,8 +157,56 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test @view(a[Block(1, 1)]) == a[Block(1, 1)]
     end
 
+    @testset "GradedOneTo" begin
+      r = gradedrange([U1(0) => 2, U1(1) => 2])
+      a = BlockSparseArray{elt}(r, r)
+      @views for i in [Block(1, 1), Block(2, 2)]
+        a[i] = randn(elt, size(a[i]))
+      end
+      b = 2 * a
+      @test block_nstored(b) == 2
+      @test Array(b) == 2 * Array(a)
+      for i in 1:2
+        @test axes(b, i) isa GradedOneTo
+        @test axes(a[:, :], i) isa GradedOneTo
+      end
+
+      I = [Block(1)[1:1]]
+      @test a[I, :] isa AbstractBlockArray
+      @test a[:, I] isa AbstractBlockArray
+      @test size(a[I, I]) == (1, 1)
+      @test !isdual(axes(a[I, I], 1))
+    end
+
+    @testset "GradedUnitRange" begin
+      r = gradedrange([U1(0) => 2, U1(1) => 2])[1:3]
+      a = BlockSparseArray{elt}(r, r)
+      @views for i in [Block(1, 1), Block(2, 2)]
+        a[i] = randn(elt, size(a[i]))
+      end
+      b = 2 * a
+      @test block_nstored(b) == 2
+      @test Array(b) == 2 * Array(a)
+      for i in 1:2
+        @test axes(b, i) isa GradedUnitRange
+        @test axes(a[:, :], i) isa GradedUnitRange
+      end
+
+      I = [Block(1)[1:1]]
+      @test a[I, :] isa AbstractBlockArray
+      @test axes(a[I, :], 1) isa GradedOneTo
+      @test axes(a[I, :], 2) isa GradedUnitRange
+
+      @test a[:, I] isa AbstractBlockArray
+      @test axes(a[:, I], 2) isa GradedOneTo
+      @test axes(a[:, I], 1) isa GradedUnitRange
+      @test size(a[I, I]) == (1, 1)
+      @test !isdual(axes(a[I, I], 1))
+    end
+
     # Test case when all axes are dual.
-    for r in (gradedrange([U1(0) => 2, U1(1) => 2]), blockedrange([2, 2]))
+    @testset "dual GradedOneTo" begin
+      r = gradedrange([U1(-1) => 2, U1(1) => 2])
       a = BlockSparseArray{elt}(dual(r), dual(r))
       @views for i in [Block(1, 1), Block(2, 2)]
         a[i] = randn(elt, size(a[i]))
@@ -157,14 +214,75 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       b = 2 * a
       @test block_nstored(b) == 2
       @test Array(b) == 2 * Array(a)
-      for ax in axes(b)
-        @test ax isa UnitRangeDual
+      for i in 1:2
+        @test axes(b, i) isa GradedUnitRangeDual
+        @test axes(a[:, :], i) isa GradedUnitRangeDual
       end
+      I = [Block(1)[1:1]]
+      @test a[I, :] isa AbstractBlockArray
+      @test a[:, I] isa AbstractBlockArray
+      @test size(a[I, I]) == (1, 1)
+      @test isdual(axes(a[I, :], 2))
+      @test isdual(axes(a[:, I], 1))
+      @test_broken isdual(axes(a[I, :], 1))
+      @test_broken isdual(axes(a[:, I], 2))
+      @test_broken isdual(axes(a[I, I], 1))
+      @test_broken isdual(axes(a[I, I], 2))
+    end
+
+    @testset "dual GradedUnitRange" begin
+      r = gradedrange([U1(0) => 2, U1(1) => 2])[1:3]
+      a = BlockSparseArray{elt}(dual(r), dual(r))
+      @views for i in [Block(1, 1), Block(2, 2)]
+        a[i] = randn(elt, size(a[i]))
+      end
+      b = 2 * a
+      @test block_nstored(b) == 2
+      @test Array(b) == 2 * Array(a)
+      for i in 1:2
+        @test axes(b, i) isa GradedUnitRangeDual
+        @test axes(a[:, :], i) isa GradedUnitRangeDual
+      end
+
+      I = [Block(1)[1:1]]
+      @test a[I, :] isa AbstractBlockArray
+      @test a[:, I] isa AbstractBlockArray
+      @test size(a[I, I]) == (1, 1)
+      @test isdual(axes(a[I, :], 2))
+      @test isdual(axes(a[:, I], 1))
+      @test_broken isdual(axes(a[I, :], 1))
+      @test_broken isdual(axes(a[:, I], 2))
+      @test_broken isdual(axes(a[I, I], 1))
+      @test_broken isdual(axes(a[I, I], 2))
     end
 
-    # Test case when all axes are dual
-    # from taking the adjoint.
-    for r in (gradedrange([U1(0) => 2, U1(1) => 2]), blockedrange([2, 2]))
+    @testset "dual BlockedUnitRange" begin    # self dual
+      r = blockedrange([2, 2])
+      a = BlockSparseArray{elt}(dual(r), dual(r))
+      @views for i in [Block(1, 1), Block(2, 2)]
+        a[i] = randn(elt, size(a[i]))
+      end
+      b = 2 * a
+      @test block_nstored(b) == 2
+      @test Array(b) == 2 * Array(a)
+      @test a[:, :] isa BlockSparseArray
+      for i in 1:2
+        @test axes(b, i) isa BlockedOneTo
+        @test axes(a[:, :], i) isa BlockedOneTo
+      end
+
+      I = [Block(1)[1:1]]
+      @test a[I, :] isa BlockSparseArray
+      @test a[:, I] isa BlockSparseArray
+      @test size(a[I, I]) == (1, 1)
+      @test !isdual(axes(a[I, I], 1))
+    end
+
+    # Test case when all axes are dual from taking the adjoint.
+    for r in (
+      gradedrange([U1(0) => 2, U1(1) => 2]),
+      gradedrange([U1(0) => 2, U1(1) => 2])[begin:end],
+    )
       a = BlockSparseArray{elt}(r, r)
       @views for i in [Block(1, 1), Block(2, 2)]
         a[i] = randn(elt, size(a[i]))
@@ -173,8 +291,13 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test block_nstored(b) == 2
       @test Array(b) == 2 * Array(a)'
       for ax in axes(b)
-        @test ax isa UnitRangeDual
+        @test ax isa typeof(dual(r))
       end
+
+      I = [Block(1)[1:1]]
+      @test size(b[I, :]) == (1, 4)
+      @test size(b[:, I]) == (4, 1)
+      @test size(b[I, I]) == (1, 1)
     end
   end
   @testset "Matrix multiplication" begin
diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl
index e409ed5500..aa3c7711c6 100644
--- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl
+++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl
@@ -1,5 +1,12 @@
 using BlockArrays:
-  BlockArrays, Block, BlockIndexRange, BlockedVector, blocklength, blocksize, viewblock
+  AbstractBlockedUnitRange,
+  BlockArrays,
+  Block,
+  BlockIndexRange,
+  BlockedVector,
+  blocklength,
+  blocksize,
+  viewblock
 
 # This splits `BlockIndexRange{N}` into
 # `NTuple{N,BlockIndexRange{1}}`.
@@ -191,7 +198,9 @@ function to_blockindexrange(
   # work right now.
   return blocks(a.blocks)[Int(I)]
 end
-function to_blockindexrange(a::Base.Slice{<:BlockedOneTo{<:Integer}}, I::Block{1})
+function to_blockindexrange(
+  a::Base.Slice{<:AbstractBlockedUnitRange{<:Integer}}, I::Block{1}
+)
   @assert I in only(blockaxes(a.indices))
   return I
 end
diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl
index a7797766d4..095364687f 100644
--- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl
+++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl
@@ -15,9 +15,15 @@ using BlockArrays:
   blocksizes,
   mortar
 using Compat: @compat
-using LinearAlgebra: mul!
+using LinearAlgebra: Adjoint, mul!
 using NDTensors.BlockSparseArrays:
-  @view!, BlockSparseArray, BlockView, block_nstored, block_reshape, view!
+  @view!,
+  BlockSparseArray,
+  BlockView,
+  block_nstored,
+  block_reshape,
+  block_stored_indices,
+  view!
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: contract
 using Test: @test, @test_broken, @test_throws, @testset
@@ -44,6 +50,17 @@ include("TestBlockSparseArraysUtils.jl")
     a[Block(2, 2)] = randn(elt, 3, 3)
     @test a[2:4, 4] == Array(a)[2:4, 4]
     @test_broken a[4, 2:4]
+
+    @test a[Block(1), :] isa BlockSparseArray{elt}
+    @test adjoint(a) isa Adjoint{elt,<:BlockSparseArray}
+    @test_broken adjoint(a)[Block(1), :] isa Adjoint{elt,<:BlockSparseArray}
+    # could also be directly a BlockSparseArray
+
+    a = BlockSparseArray{elt}([1], [1, 1])
+    a[1, 2] = 1
+    @test [a[Block(Tuple(it))] for it in eachindex(block_stored_indices(a))] isa Vector
+    ah = adjoint(a)
+    @test_broken [ah[Block(Tuple(it))] for it in eachindex(block_stored_indices(ah))] isa Vector
   end
   @testset "Basics" begin
     a = BlockSparseArray{elt}([2, 3], [2, 3])
diff --git a/NDTensors/src/lib/GradedAxes/src/GradedAxes.jl b/NDTensors/src/lib/GradedAxes/src/GradedAxes.jl
index d605289f81..7edd09bf84 100644
--- a/NDTensors/src/lib/GradedAxes/src/GradedAxes.jl
+++ b/NDTensors/src/lib/GradedAxes/src/GradedAxes.jl
@@ -2,6 +2,7 @@ module GradedAxes
 include("blockedunitrange.jl")
 include("gradedunitrange.jl")
 include("dual.jl")
-include("unitrangedual.jl")
+include("gradedunitrangedual.jl")
+include("onetoone.jl")
 include("fusion.jl")
 end
diff --git a/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl b/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl
index 883025df12..d913dd60ab 100644
--- a/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl
+++ b/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl
@@ -167,7 +167,7 @@ end
 # Slice `a` by `I`, returning a:
 # `BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}`
 # with the `BlockIndex{1}` corresponding to each value of `I`.
-function to_blockindices(a::BlockedOneTo{<:Integer}, I::UnitRange{<:Integer})
+function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::UnitRange{<:Integer})
   return mortar(
     map(blocks(blockedunitrange_getindices(a, I))) do r
       bi_first = findblockindex(a, first(r))
diff --git a/NDTensors/src/lib/GradedAxes/src/dual.jl b/NDTensors/src/lib/GradedAxes/src/dual.jl
index ff11b1103c..ca985e30a0 100644
--- a/NDTensors/src/lib/GradedAxes/src/dual.jl
+++ b/NDTensors/src/lib/GradedAxes/src/dual.jl
@@ -1,7 +1,12 @@
-function dual end
+# default behavior: self-dual
+dual(r::AbstractUnitRange) = r
+nondual(r::AbstractUnitRange) = r
+isdual(::AbstractUnitRange) = false
 
 using NDTensors.LabelledNumbers:
   LabelledStyle, IsLabelled, NotLabelled, label, labelled, unlabel
+
+dual(i::LabelledInteger) = labelled(unlabel(i), dual(label(i)))
 label_dual(x) = label_dual(LabelledStyle(x), x)
 label_dual(::NotLabelled, x) = x
 label_dual(::IsLabelled, x) = labelled(unlabel(x), dual(label(x)))
diff --git a/NDTensors/src/lib/GradedAxes/src/fusion.jl b/NDTensors/src/lib/GradedAxes/src/fusion.jl
index 3c3b58bc91..2e893ec1db 100644
--- a/NDTensors/src/lib/GradedAxes/src/fusion.jl
+++ b/NDTensors/src/lib/GradedAxes/src/fusion.jl
@@ -1,12 +1,5 @@
 using BlockArrays: AbstractBlockedUnitRange, blocklengths
 
-# Represents the range `1:1` or `Base.OneTo(1)`.
-struct OneToOne{T} <: AbstractUnitRange{T} end
-OneToOne() = OneToOne{Bool}()
-Base.first(a::OneToOne) = one(eltype(a))
-Base.last(a::OneToOne) = one(eltype(a))
-BlockArrays.blockaxes(g::OneToOne) = (Block.(g),)  # BlockArrays default crashes for OneToOne{Bool}
-
 # https://github.com/ITensor/ITensors.jl/blob/v0.3.57/NDTensors/src/lib/GradedAxes/src/tensor_product.jl
 # https://en.wikipedia.org/wiki/Tensor_product
 # https://github.com/KeitaNakamura/Tensorial.jl
@@ -20,7 +13,7 @@ function tensor_product(
 end
 
 flip_dual(r::AbstractUnitRange) = r
-flip_dual(r::UnitRangeDual) = flip(r)
+flip_dual(r::GradedUnitRangeDual) = flip(r)
 function tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange)
   return tensor_product(flip_dual(a1), flip_dual(a2))
 end
@@ -67,7 +60,7 @@ function tensor_product(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRan
   return blockedrange(new_blocklengths)
 end
 
-# convention: sort UnitRangeDual according to nondual blocks
+# convention: sort GradedUnitRangeDual according to nondual blocks
 function blocksortperm(a::AbstractUnitRange)
   return Block.(sortperm(blocklabels(nondual(a))))
 end
@@ -102,7 +95,7 @@ function blockmergesort(g::AbstractGradedUnitRange)
   return gradedrange(new_blocklengths)
 end
 
-blockmergesort(g::UnitRangeDual) = flip(blockmergesort(flip(g)))
+blockmergesort(g::GradedUnitRangeDual) = flip(blockmergesort(flip(g)))
 blockmergesort(g::AbstractUnitRange) = g
 
 # fusion_product produces a sorted, non-dual GradedUnitRange
@@ -111,7 +104,7 @@ function fusion_product(g1, g2)
 end
 
 fusion_product(g::AbstractUnitRange) = blockmergesort(g)
-fusion_product(g::UnitRangeDual) = fusion_product(flip(g))
+fusion_product(g::GradedUnitRangeDual) = fusion_product(flip(g))
 
 # recursive fusion_product. Simpler than reduce + fix type stability issues with reduce
 function fusion_product(g1, g2, g3...)
diff --git a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
index 08fcef5e89..0bd35707a7 100644
--- a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
+++ b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
@@ -17,7 +17,8 @@ using BlockArrays:
   blocklengths,
   findblock,
   findblockindex,
-  mortar
+  mortar,
+  sortedunion
 using Compat: allequal
 using FillArrays: Fill
 using ..LabelledNumbers:
@@ -25,23 +26,43 @@ using ..LabelledNumbers:
   LabelledInteger,
   LabelledUnitRange,
   label,
+  label_type,
   labelled,
   labelled_isequal,
   unlabel
 
-const AbstractGradedUnitRange{T<:LabelledInteger} = AbstractBlockedUnitRange{T}
+abstract type AbstractGradedUnitRange{T,BlockLasts} <:
+              AbstractBlockedUnitRange{T,BlockLasts} end
 
-const GradedUnitRange{T<:LabelledInteger,BlockLasts<:Vector{T}} = BlockedUnitRange{
-  T,BlockLasts
-}
+struct GradedUnitRange{T,BlockLasts<:Vector{T}} <: AbstractGradedUnitRange{T,BlockLasts}
+  first::T
+  lasts::BlockLasts
+end
 
-const GradedOneTo{T<:LabelledInteger,BlockLasts<:Vector{T}} = BlockedOneTo{T,BlockLasts}
+struct GradedOneTo{T,BlockLasts<:Vector{T}} <: AbstractGradedUnitRange{T,BlockLasts}
+  lasts::BlockLasts
 
-# This is only needed in certain Julia versions below 1.10
-# (for example Julia 1.6).
-# TODO: Delete this once we drop Julia 1.6 support.
-function Base.OrdinalRange{T,T}(a::GradedOneTo{<:LabelledInteger{T}}) where {T}
-  return unlabel_blocks(a)
+  # assume that lasts is sorted, no checks carried out here
+  function GradedOneTo(lasts::BlockLasts) where {T<:Integer,BlockLasts<:AbstractVector{T}}
+    Base.require_one_based_indexing(lasts)
+    isempty(lasts) || first(lasts) >= 0 || throw(ArgumentError("blocklasts must be >= 0"))
+    return new{T,BlockLasts}(lasts)
+  end
+  function GradedOneTo(lasts::BlockLasts) where {T<:Integer,BlockLasts<:Tuple{T,Vararg{T}}}
+    first(lasts) >= 0 || throw(ArgumentError("blocklasts must be >= 0"))
+    return new{T,BlockLasts}(lasts)
+  end
+end
+
+function Base.show(io::IO, ::MIME"text/plain", g::AbstractGradedUnitRange)
+  v = map(b -> label(b) => unlabel(b), blocks(g))
+  println(io, typeof(g))
+  return print(io, join(repr.(v), '\n'))
+end
+
+function Base.show(io::IO, g::AbstractGradedUnitRange)
+  v = map(b -> label(b) => unlabel(b), blocks(g))
+  return print(io, nameof(typeof(g)), '[', join(repr.(v), ", "), ']')
 end
 
 # == is just a range comparison that ignores labels. Need dedicated function to check equality.
@@ -56,41 +77,22 @@ function space_isequal(a1::AbstractUnitRange, a2::AbstractUnitRange)
   return (isdual(a1) == isdual(a2)) && labelled_isequal(a1, a2)
 end
 
-# This is only needed in certain Julia versions below 1.10
-# (for example Julia 1.6).
-# TODO: Delete this once we drop Julia 1.6 support.
-# The type constraint `T<:Integer` is needed to avoid an ambiguity
-# error with a conversion method in Base.
-function Base.UnitRange{T}(
+# needed in BlockSparseArrays
+function Base.AbstractUnitRange{T}(
   a::AbstractGradedUnitRange{<:LabelledInteger{T}}
-) where {T<:Integer}
-  return UnitRange(unlabel_blocks(a))
-end
-
-# This is only needed in certain Julia versions below 1.10
-# (for example Julia 1.6).
-# TODO: Delete this once we drop Julia 1.6 support.
-# The type constraint `T<:Integer` is needed to avoid an ambiguity
-# error with a conversion method in Base.
-using BlockArrays: BlockSlice
-function Base.UnitRange{T}(
-  a::BlockSlice{<:Any,<:LabelledInteger{T},<:AbstractUnitRange{<:LabelledInteger{T}}}
-) where {T<:Integer}
-  return UnitRange{T}(a.indices)
-end
-
-# TODO: See if this is needed.
-function Base.AbstractUnitRange{T}(a::GradedOneTo{<:LabelledInteger{T}}) where {T}
+) where {T}
   return unlabel_blocks(a)
 end
 
 # TODO: Use `TypeParameterAccessors`.
 Base.eltype(::Type{<:GradedUnitRange{T}}) where {T} = T
+LabelledNumbers.label_type(g::AbstractGradedUnitRange) = label_type(typeof(g))
+LabelledNumbers.label_type(T::Type{<:AbstractGradedUnitRange}) = label_type(eltype(T))
 
 function gradedrange(lblocklengths::AbstractVector{<:LabelledInteger})
   brange = blockedrange(unlabel.(lblocklengths))
   lblocklasts = labelled.(blocklasts(brange), label.(lblocklengths))
-  return BlockedOneTo(lblocklasts)
+  return GradedOneTo(lblocklasts)
 end
 
 # To help with generic code.
@@ -100,17 +102,6 @@ end
 
 Base.last(a::AbstractGradedUnitRange) = isempty(a.lasts) ? first(a) - 1 : last(a.lasts)
 
-# TODO: This needs to be defined to circumvent an issue
-# in the `BlockArrays.BlocksView` constructor. This
-# is likely caused by issues around `BlockedUnitRange` constraining
-# the element type to be `Int`, which is being fixed in:
-# https://github.com/JuliaArrays/BlockArrays.jl/pull/337
-# Remove this definition once that is fixed.
-function BlockArrays.blocks(a::AbstractGradedUnitRange)
-  # TODO: Fix `BlockRange`, try using `BlockRange` instead.
-  return [a[Block(i)] for i in 1:blocklength(a)]
-end
-
 function gradedrange(lblocklengths::AbstractVector{<:Pair{<:Any,<:Integer}})
   return gradedrange(labelled.(last.(lblocklengths), first.(lblocklengths)))
 end
@@ -118,14 +109,12 @@ end
 function labelled_blocks(a::BlockedOneTo, labels)
   # TODO: Use `blocklasts(a)`? That might
   # cause a recursive loop.
-  return BlockedOneTo(labelled.(a.lasts, labels))
+  return GradedOneTo(labelled.(a.lasts, labels))
 end
 function labelled_blocks(a::BlockedUnitRange, labels)
   # TODO: Use `first(a)` and `blocklasts(a)`? Those might
   # cause a recursive loop.
-  return BlockArrays._BlockedUnitRange(
-    labelled(a.first, labels[1]), labelled.(a.lasts, labels)
-  )
+  return GradedUnitRange(labelled(a.first, labels[1]), labelled.(a.lasts, labels))
 end
 
 function BlockArrays.findblock(a::AbstractGradedUnitRange, index::Integer)
@@ -177,15 +166,15 @@ end
 # TODO: This relies on internals of `BlockArrays`, maybe redesign
 # to try to avoid that.
 # TODO: Define `set_grades`, `set_sector_labels`, `set_labels`.
-function unlabel_blocks(a::BlockedOneTo)
+function unlabel_blocks(a::GradedOneTo)
   # TODO: Use `blocklasts(a)`.
   return BlockedOneTo(unlabel.(a.lasts))
 end
-function unlabel_blocks(a::BlockedUnitRange)
+function unlabel_blocks(a::GradedUnitRange)
   return BlockArrays._BlockedUnitRange(a.first, unlabel.(a.lasts))
 end
 
-## BlockedUnitRage interface
+## BlockedUnitRange interface
 
 function Base.axes(ga::AbstractGradedUnitRange)
   return map(axes(unlabel_blocks(ga))) do a
@@ -199,9 +188,6 @@ end
 function BlockArrays.blockfirsts(a::AbstractGradedUnitRange)
   return gradedunitrange_blockfirsts(a)
 end
-function BlockArrays.blockfirsts(a::GradedOneTo)
-  return gradedunitrange_blockfirsts(a)
-end
 
 function BlockArrays.blocklasts(a::AbstractGradedUnitRange)
   return labelled.(blocklasts(unlabel_blocks(a)), blocklabels(a))
@@ -217,9 +203,6 @@ end
 function Base.first(a::AbstractGradedUnitRange)
   return gradedunitrange_first(a)
 end
-function Base.first(a::GradedOneTo)
-  return gradedunitrange_first(a)
-end
 
 Base.iterate(a::AbstractGradedUnitRange) = isempty(a) ? nothing : (first(a), first(a))
 function Base.iterate(a::AbstractGradedUnitRange, i)
@@ -232,12 +215,44 @@ function firstblockindices(a::AbstractGradedUnitRange)
   return labelled.(firstblockindices(unlabel_blocks(a)), blocklabels(a))
 end
 
-function blockedunitrange_getindex(a::AbstractGradedUnitRange, index)
-  # This uses `blocklasts` since that is what is stored
-  # in `BlockedUnitRange`, maybe abstract that away.
+function blockedunitrange_getindices(a::AbstractGradedUnitRange, index::Block{1})
+  return labelled(unlabel_blocks(a)[index], get_label(a, index))
+end
+
+function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::Vector{<:Integer})
+  return map(index -> a[index], indices)
+end
+
+function blockedunitrange_getindices(
+  a::AbstractGradedUnitRange,
+  indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+
+function blockedunitrange_getindices(a::AbstractGradedUnitRange, index)
   return labelled(unlabel_blocks(a)[index], get_label(a, index))
 end
 
+function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::BlockIndexRange)
+  return a[block(indices)][only(indices.indices)]
+end
+
+function blockedunitrange_getindices(
+  a::AbstractGradedUnitRange, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}}
+)
+  # Without converting `indices` to `Vector`,
+  # mapping `indices` outputs a `BlockVector`
+  # which is harder to reason about.
+  blocks = map(index -> a[index], Vector(indices))
+  # We pass `length.(blocks)` to `mortar` in order
+  # to pass block labels to the axes of the output,
+  # if they exist. This makes it so that
+  # `only(axes(a[indices])) isa `GradedUnitRange`
+  # if `a isa `GradedUnitRange`, for example.
+  return mortar(blocks, length.(blocks))
+end
+
 # The block labels of the corresponding slice.
 function blocklabels(a::AbstractUnitRange, indices)
   return map(_blocks(a, indices)) do block
@@ -252,11 +267,6 @@ function blockedunitrange_getindices(
   return labelled_blocks(a_indices, blocklabels(ga, indices))
 end
 
-# Fixes ambiguity error with:
-# ```julia
-# blockedunitrange_getindices(::GradedUnitRange, ::AbstractUnitRange{<:Integer})
-# ```
-# TODO: Try removing once GradedAxes is rewritten for BlockArrays v1.
 function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::BlockSlice)
   return a[indices.block]
 end
@@ -270,28 +280,25 @@ function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::BlockI
 end
 
 function Base.getindex(a::AbstractGradedUnitRange, index::Integer)
-  # This uses `blocklasts` since that is what is stored
-  # in `BlockedUnitRange`, maybe abstract that away.
   return labelled(unlabel_blocks(a)[index], get_label(a, index))
 end
 
 function Base.getindex(a::AbstractGradedUnitRange, index::Block{1})
-  return blockedunitrange_getindex(a, index)
+  return blockedunitrange_getindices(a, index)
 end
 
 function Base.getindex(a::AbstractGradedUnitRange, indices::BlockIndexRange)
   return blockedunitrange_getindices(a, indices)
 end
 
+# fix ambiguities
 function Base.getindex(
-  a::AbstractGradedUnitRange, indices::BlockRange{1,<:Tuple{AbstractUnitRange{Int}}}
+  a::AbstractGradedUnitRange, indices::BlockArrays.BlockRange{1,<:Tuple{Base.OneTo}}
 )
   return blockedunitrange_getindices(a, indices)
 end
-
-# Fixes ambiguity error with `BlockArrays`.
 function Base.getindex(
-  a::AbstractGradedUnitRange, indices::BlockRange{1,Tuple{Base.OneTo{Int}}}
+  a::AbstractGradedUnitRange, indices::BlockRange{1,<:Tuple{AbstractUnitRange{Int}}}
 )
   return blockedunitrange_getindices(a, indices)
 end
@@ -307,8 +314,6 @@ end
 # getindex(::GradedUnitRange, ::Any)
 # getindex(::AbstractUnitRange, ::AbstractUnitRange{<:Integer})
 # ```
-# TODO: Maybe not needed once GradedAxes is rewritten
-# for BlockArrays v1.
 function Base.getindex(a::AbstractGradedUnitRange, indices::BlockSlice)
   return blockedunitrange_getindices(a, indices)
 end
@@ -326,17 +331,30 @@ end
 # that mixed dense and graded axes.
 # TODO: Maybe come up with a more general solution.
 function BlockArrays.combine_blockaxes(
-  a1::GradedOneTo{<:LabelledInteger{T}}, a2::Base.OneTo{T}
+  a1::AbstractGradedUnitRange{<:LabelledInteger{T}}, a2::AbstractUnitRange{T}
 ) where {T<:Integer}
   combined_blocklasts = sort!(union(unlabel.(blocklasts(a1)), blocklasts(a2)))
   return BlockedOneTo(combined_blocklasts)
 end
 function BlockArrays.combine_blockaxes(
-  a1::Base.OneTo{T}, a2::GradedOneTo{<:LabelledInteger{T}}
+  a1::AbstractUnitRange{T}, a2::AbstractGradedUnitRange{<:LabelledInteger{T}}
 ) where {T<:Integer}
   return BlockArrays.combine_blockaxes(a2, a1)
 end
 
+# preserve labels inside combine_blockaxes
+function BlockArrays.combine_blockaxes(a::GradedOneTo, b::GradedOneTo)
+  return GradedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
+end
+function BlockArrays.combine_blockaxes(a::GradedUnitRange, b::GradedUnitRange)
+  new_blocklasts = sortedunion(blocklasts(a), blocklasts(b))
+  new_first = labelled(oneunit(eltype(new_blocklasts)), label(first(new_blocklasts)))
+  return GradedUnitRange(new_first, new_blocklasts)
+end
+
+# preserve axes in SubArray
+Base.axes(S::Base.Slice{<:AbstractGradedUnitRange}) = (S.indices,)
+
 # Version of length that checks that all blocks have the same label
 # and returns a labelled length with that label.
 function labelled_length(a::AbstractBlockVector{<:Integer})
diff --git a/NDTensors/src/lib/GradedAxes/src/gradedunitrangedual.jl b/NDTensors/src/lib/GradedAxes/src/gradedunitrangedual.jl
new file mode 100644
index 0000000000..217d4b401f
--- /dev/null
+++ b/NDTensors/src/lib/GradedAxes/src/gradedunitrangedual.jl
@@ -0,0 +1,105 @@
+struct GradedUnitRangeDual{
+  T,BlockLasts,NondualUnitRange<:AbstractGradedUnitRange{T,BlockLasts}
+} <: AbstractGradedUnitRange{T,BlockLasts}
+  nondual_unitrange::NondualUnitRange
+end
+
+dual(a::AbstractGradedUnitRange) = GradedUnitRangeDual(a)
+nondual(a::GradedUnitRangeDual) = a.nondual_unitrange
+dual(a::GradedUnitRangeDual) = nondual(a)
+flip(a::GradedUnitRangeDual) = dual(flip(nondual(a)))
+isdual(::GradedUnitRangeDual) = true
+## TODO: Define this to instantiate a dual unit range.
+## materialize_dual(a::GradedUnitRangeDual) = materialize_dual(nondual(a))
+
+Base.first(a::GradedUnitRangeDual) = label_dual(first(nondual(a)))
+Base.last(a::GradedUnitRangeDual) = label_dual(last(nondual(a)))
+Base.step(a::GradedUnitRangeDual) = label_dual(step(nondual(a)))
+
+Base.view(a::GradedUnitRangeDual, index::Block{1}) = a[index]
+
+function blockedunitrange_getindices(
+  a::GradedUnitRangeDual, indices::AbstractUnitRange{<:Integer}
+)
+  return dual(getindex(nondual(a), indices))
+end
+
+using BlockArrays: Block, BlockIndexRange, BlockRange
+
+function blockedunitrange_getindices(a::GradedUnitRangeDual, indices::Integer)
+  return label_dual(getindex(nondual(a), indices))
+end
+
+function blockedunitrange_getindices(a::GradedUnitRangeDual, indices::Block{1})
+  return label_dual(getindex(nondual(a), indices))
+end
+
+function blockedunitrange_getindices(a::GradedUnitRangeDual, indices::BlockRange)
+  return label_dual(getindex(nondual(a), indices))
+end
+
+# fix ambiguity
+function blockedunitrange_getindices(
+  a::GradedUnitRangeDual, indices::BlockRange{1,<:Tuple{AbstractUnitRange{Int}}}
+)
+  return dual(getindex(nondual(a), indices))
+end
+
+function BlockArrays.blocklengths(a::GradedUnitRangeDual)
+  return dual.(blocklengths(nondual(a)))
+end
+
+function gradedunitrangedual_getindices_blocks(a::GradedUnitRangeDual, indices)
+  a_indices = getindex(nondual(a), indices)
+  return mortar([label_dual(b) for b in blocks(a_indices)])
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::GradedUnitRangeDual, indices::Vector{<:Block{1}})
+  return gradedunitrangedual_getindices_blocks(a, indices)
+end
+
+function blockedunitrange_getindices(
+  a::GradedUnitRangeDual, indices::Vector{<:BlockIndexRange{1}}
+)
+  return gradedunitrangedual_getindices_blocks(a, indices)
+end
+
+Base.axes(a::GradedUnitRangeDual) = axes(nondual(a))
+
+using BlockArrays: BlockArrays, Block, BlockSlice
+using NDTensors.LabelledNumbers: LabelledUnitRange
+function BlockArrays.BlockSlice(b::Block, a::LabelledUnitRange)
+  return BlockSlice(b, unlabel(a))
+end
+
+using BlockArrays: BlockArrays, BlockSlice
+using NDTensors.GradedAxes: GradedUnitRangeDual, dual
+function BlockArrays.BlockSlice(b::Block, r::GradedUnitRangeDual)
+  return BlockSlice(b, dual(r))
+end
+
+using NDTensors.LabelledNumbers: LabelledNumbers, LabelledUnitRange, label
+function Base.iterate(a::GradedUnitRangeDual, i)
+  i == last(a) && return nothing
+  return dual.(iterate(nondual(a), i))
+end
+
+using NDTensors.LabelledNumbers: LabelledInteger, label, labelled, unlabel
+using BlockArrays: BlockArrays, blockaxes, blocklasts, combine_blockaxes, findblock
+BlockArrays.blockaxes(a::GradedUnitRangeDual) = blockaxes(nondual(a))
+BlockArrays.blockfirsts(a::GradedUnitRangeDual) = label_dual.(blockfirsts(nondual(a)))
+BlockArrays.blocklasts(a::GradedUnitRangeDual) = label_dual.(blocklasts(nondual(a)))
+function BlockArrays.findblock(a::GradedUnitRangeDual, index::Integer)
+  return findblock(nondual(a), index)
+end
+
+blocklabels(a::GradedUnitRangeDual) = dual.(blocklabels(nondual(a)))
+
+function BlockArrays.combine_blockaxes(a1::GradedUnitRangeDual, a2::GradedUnitRangeDual)
+  return dual(combine_blockaxes(nondual(a1), nondual(a2)))
+end
+
+function unlabel_blocks(a::GradedUnitRangeDual)
+  return unlabel_blocks(nondual(a))
+end
diff --git a/NDTensors/src/lib/GradedAxes/src/onetoone.jl b/NDTensors/src/lib/GradedAxes/src/onetoone.jl
new file mode 100644
index 0000000000..426df396b1
--- /dev/null
+++ b/NDTensors/src/lib/GradedAxes/src/onetoone.jl
@@ -0,0 +1,9 @@
+using BlockArrays: AbstractBlockedUnitRange
+using ..LabelledNumbers: islabelled
+
+# Represents the range `1:1` or `Base.OneTo(1)`.
+struct OneToOne{T} <: AbstractUnitRange{T} end
+OneToOne() = OneToOne{Bool}()
+Base.first(a::OneToOne) = one(eltype(a))
+Base.last(a::OneToOne) = one(eltype(a))
+BlockArrays.blockaxes(g::OneToOne) = (Block.(g),) # BlockArrays default crashes for OneToOne{Bool}
diff --git a/NDTensors/src/lib/GradedAxes/src/unitrangedual.jl b/NDTensors/src/lib/GradedAxes/src/unitrangedual.jl
deleted file mode 100644
index 8e9529e2d0..0000000000
--- a/NDTensors/src/lib/GradedAxes/src/unitrangedual.jl
+++ /dev/null
@@ -1,135 +0,0 @@
-struct UnitRangeDual{T,NondualUnitRange<:AbstractUnitRange} <: AbstractUnitRange{T}
-  nondual_unitrange::NondualUnitRange
-end
-UnitRangeDual(a::AbstractUnitRange) = UnitRangeDual{eltype(a),typeof(a)}(a)
-
-dual(a::AbstractUnitRange) = UnitRangeDual(a)
-nondual(a::UnitRangeDual) = a.nondual_unitrange
-dual(a::UnitRangeDual) = nondual(a)
-flip(a::UnitRangeDual) = dual(flip(nondual(a)))
-nondual(a::AbstractUnitRange) = a
-isdual(::AbstractUnitRange) = false
-isdual(::UnitRangeDual) = true
-## TODO: Define this to instantiate a dual unit range.
-## materialize_dual(a::UnitRangeDual) = materialize_dual(nondual(a))
-
-Base.first(a::UnitRangeDual) = label_dual(first(nondual(a)))
-Base.last(a::UnitRangeDual) = label_dual(last(nondual(a)))
-Base.step(a::UnitRangeDual) = label_dual(step(nondual(a)))
-
-Base.view(a::UnitRangeDual, index::Block{1}) = a[index]
-
-function Base.show(io::IO, a::UnitRangeDual)
-  return print(io, UnitRangeDual, "(", blocklasts(a), ")")
-end
-
-function Base.show(io::IO, mimetype::MIME"text/plain", a::UnitRangeDual)
-  return Base.invoke(
-    show, Tuple{typeof(io),MIME"text/plain",AbstractArray}, io, mimetype, a
-  )
-end
-
-function Base.getindex(a::UnitRangeDual, indices::AbstractUnitRange{<:Integer})
-  return dual(getindex(nondual(a), indices))
-end
-
-using BlockArrays: Block, BlockIndexRange, BlockRange
-
-function Base.getindex(a::UnitRangeDual, indices::Integer)
-  return label_dual(getindex(nondual(a), indices))
-end
-
-# TODO: Use `label_dual.` here, make broadcasting work?
-Base.getindex(a::UnitRangeDual, indices::Block{1}) = dual(getindex(nondual(a), indices))
-
-# TODO: Use `label_dual.` here, make broadcasting work?
-Base.getindex(a::UnitRangeDual, indices::BlockRange) = dual(getindex(nondual(a), indices))
-
-# TODO: Use `label_dual.` here, make broadcasting work?
-function unitrangedual_getindices_blocks(a, indices)
-  a_indices = getindex(nondual(a), indices)
-  return mortar([dual(b) for b in blocks(a_indices)])
-end
-
-# TODO: Move this to a `BlockArraysExtensions` library.
-function blockedunitrange_getindices(a::UnitRangeDual, indices::Block{1})
-  return a[indices]
-end
-
-function Base.getindex(a::UnitRangeDual, indices::Vector{<:Block{1}})
-  return unitrangedual_getindices_blocks(a, indices)
-end
-
-function Base.getindex(a::UnitRangeDual, indices::Vector{<:BlockIndexRange{1}})
-  return unitrangedual_getindices_blocks(a, indices)
-end
-
-function to_blockindices(a::UnitRangeDual, indices::UnitRange{<:Integer})
-  return to_blockindices(nondual(a), indices)
-end
-
-Base.axes(a::UnitRangeDual) = axes(nondual(a))
-
-using BlockArrays: BlockArrays, Block, BlockSlice
-using NDTensors.LabelledNumbers: LabelledUnitRange
-function BlockArrays.BlockSlice(b::Block, a::LabelledUnitRange)
-  return BlockSlice(b, unlabel(a))
-end
-
-using BlockArrays: BlockArrays, BlockSlice
-using NDTensors.GradedAxes: UnitRangeDual, dual
-function BlockArrays.BlockSlice(b::Block, r::UnitRangeDual)
-  return BlockSlice(b, dual(r))
-end
-
-using NDTensors.LabelledNumbers: LabelledNumbers, label
-LabelledNumbers.label(a::UnitRangeDual) = dual(label(nondual(a)))
-
-using NDTensors.LabelledNumbers: LabelledUnitRange
-# The Base version of `length(::AbstractUnitRange)` drops the label.
-function Base.length(a::UnitRangeDual{<:Any,<:LabelledUnitRange})
-  return dual(length(nondual(a)))
-end
-function Base.iterate(a::UnitRangeDual, i)
-  i == last(a) && return nothing
-  return dual.(iterate(nondual(a), i))
-end
-# TODO: Is this a good definition?
-Base.unitrange(a::UnitRangeDual{<:Any,<:AbstractUnitRange}) = a
-
-using NDTensors.LabelledNumbers: LabelledInteger, label, labelled, unlabel
-dual(i::LabelledInteger) = labelled(unlabel(i), dual(label(i)))
-
-using BlockArrays: BlockArrays, blockaxes, blocklasts, combine_blockaxes, findblock
-BlockArrays.blockaxes(a::UnitRangeDual) = blockaxes(nondual(a))
-BlockArrays.blockfirsts(a::UnitRangeDual) = label_dual.(blockfirsts(nondual(a)))
-BlockArrays.blocklasts(a::UnitRangeDual) = label_dual.(blocklasts(nondual(a)))
-BlockArrays.findblock(a::UnitRangeDual, index::Integer) = findblock(nondual(a), index)
-
-blocklabels(a::UnitRangeDual) = dual.(blocklabels(nondual(a)))
-
-function BlockArrays.combine_blockaxes(a1::UnitRangeDual, a2::UnitRangeDual)
-  return dual(combine_blockaxes(dual(a1), dual(a2)))
-end
-
-# This is needed when constructing `CartesianIndices` from
-# a tuple of unit ranges that have this kind of dual unit range.
-# TODO: See if we can find some more elegant way of constructing
-# `CartesianIndices`, maybe by defining conversion of `LabelledInteger`
-# to `Int`, defining a more general `convert` function, etc.
-function Base.OrdinalRange{Int,Int}(
-  r::UnitRangeDual{<:LabelledInteger{Int},<:LabelledUnitRange{Int,UnitRange{Int}}}
-)
-  # TODO: Implement this broadcasting operation and use it here.
-  # return Int.(r)
-  return unlabel(nondual(r))
-end
-
-# This is only needed in certain Julia versions below 1.10
-# (for example Julia 1.6).
-# TODO: Delete this once we drop Julia 1.6 support.
-# The type constraint `T<:Integer` is needed to avoid an ambiguity
-# error with a conversion method in Base.
-function Base.UnitRange{T}(a::UnitRangeDual{<:LabelledInteger{T}}) where {T<:Integer}
-  return UnitRange{T}(nondual(a))
-end
diff --git a/NDTensors/src/lib/GradedAxes/test/test_basics.jl b/NDTensors/src/lib/GradedAxes/test/test_basics.jl
index a6c66c8afd..90faa59b93 100644
--- a/NDTensors/src/lib/GradedAxes/test/test_basics.jl
+++ b/NDTensors/src/lib/GradedAxes/test/test_basics.jl
@@ -1,6 +1,7 @@
 @eval module $(gensym())
 using BlockArrays:
   Block,
+  BlockRange,
   BlockSlice,
   BlockVector,
   blockedrange,
@@ -8,18 +9,42 @@ using BlockArrays:
   blocklasts,
   blocklength,
   blocklengths,
-  blocks
-using NDTensors.GradedAxes: GradedOneTo, GradedUnitRange, blocklabels, gradedrange
+  blocks,
+  combine_blockaxes,
+  mortar
+using NDTensors.GradedAxes:
+  GradedOneTo, GradedUnitRange, OneToOne, blocklabels, gradedrange, space_isequal
 using NDTensors.LabelledNumbers:
   LabelledUnitRange, islabelled, label, labelled, labelled_isequal, unlabel
 using Test: @test, @test_broken, @testset
+
+@testset "OneToOne" begin
+  a0 = OneToOne()
+  @test a0 isa OneToOne{Bool}
+  @test eltype(a0) == Bool
+  @test length(a0) == 1
+  @test labelled_isequal(a0, a0)
+  @test a0[1] == true
+  @test a0[[1]] == [true]
+
+  @test labelled_isequal(a0, 1:1)
+  @test labelled_isequal(1:1, a0)
+  @test !labelled_isequal(a0, 1:2)
+  @test !labelled_isequal(1:2, a0)
+end
+
 @testset "GradedAxes basics" begin
+  a0 = OneToOne()
   for a in (
     blockedrange([labelled(2, "x"), labelled(3, "y")]),
     gradedrange([labelled(2, "x"), labelled(3, "y")]),
     gradedrange(["x" => 2, "y" => 3]),
   )
     @test a isa GradedOneTo
+    @test labelled_isequal(a, a)
+    @test !labelled_isequal(a0, a)
+    @test !labelled_isequal(a, a0)
+    @test !labelled_isequal(a, 1:5)
     for x in iterate(a)
       @test x == 1
       @test label(x) == "x"
@@ -72,6 +97,13 @@ using Test: @test, @test_broken, @testset
     @test length(a[Block(2)]) == 3
     @test blocklengths(only(axes(a))) == blocklengths(a)
     @test blocklabels(only(axes(a))) == blocklabels(a)
+
+    @test axes(Base.Slice(a)) isa Tuple{typeof(a)}
+    @test AbstractUnitRange{Int}(a) == 1:5
+    b = combine_blockaxes(a, a)
+    @test b isa GradedOneTo
+    @test b == 1:5
+    @test space_isequal(b, a)
   end
 
   # Slicing operations
@@ -89,6 +121,15 @@ using Test: @test, @test_broken, @testset
   @test length(ax) == length(a)
   @test blocklengths(ax) == blocklengths(a)
   @test blocklabels(ax) == blocklabels(a)
+  @test blockfirsts(a) == [2, 3]
+
+  @test AbstractUnitRange{Int}(a) == 2:4
+  b = combine_blockaxes(a, a)
+  @test b isa GradedUnitRange
+  @test b == 1:4
+
+  @test x[[2, 4]] == [labelled(2, "x"), labelled(4, "y")]
+  @test labelled_isequal(x[BlockRange(1)], gradedrange(["x" => 2]))
 
   # Regression test for ambiguity error.
   x = gradedrange(["x" => 2, "y" => 3])
@@ -116,6 +157,7 @@ using Test: @test, @test_broken, @testset
   @test length(ax) == length(a)
   @test blocklengths(ax) == blocklengths(a)
   @test blocklabels(ax) == blocklabels(a)
+  @test axes(Base.Slice(a)) isa Tuple{typeof(a)}
 
   x = gradedrange(["x" => 2, "y" => 3])
   a = x[2:4][1:2]
@@ -198,5 +240,17 @@ using Test: @test, @test_broken, @testset
   # once `blocklengths(::BlockVector)` is defined.
   @test blocklengths(ax) == [2, 2]
   @test blocklabels(ax) == blocklabels(a)
+
+  x = gradedrange(["x" => 2, "y" => 3])
+  I = mortar([Block(1)[1:1]])
+  a = x[I]
+  @test length(a) == 1
+  @test label(first(a)) == "x"
+
+  x = gradedrange(["x" => 2, "y" => 3])[1:5]
+  I = mortar([Block(1)[1:1]])
+  a = x[I]
+  @test length(a) == 1
+  @test label(first(a)) == "x"
 end
 end
diff --git a/NDTensors/src/lib/GradedAxes/test/test_dual.jl b/NDTensors/src/lib/GradedAxes/test/test_dual.jl
index dc7f740713..18dbac045c 100644
--- a/NDTensors/src/lib/GradedAxes/test/test_dual.jl
+++ b/NDTensors/src/lib/GradedAxes/test/test_dual.jl
@@ -1,9 +1,23 @@
 @eval module $(gensym())
 using BlockArrays:
-  Block, blockaxes, blockfirsts, blocklasts, blocklength, blocklengths, blocks, findblock
+  Block,
+  BlockedOneTo,
+  blockaxes,
+  blockedrange,
+  blockfirsts,
+  blockisequal,
+  blocklasts,
+  blocklength,
+  blocklengths,
+  blocks,
+  findblock,
+  mortar,
+  combine_blockaxes
 using NDTensors.GradedAxes:
+  AbstractGradedUnitRange,
   GradedAxes,
-  UnitRangeDual,
+  GradedUnitRangeDual,
+  OneToOne,
   blocklabels,
   blockmergesortperm,
   blocksortperm,
@@ -13,78 +27,137 @@ using NDTensors.GradedAxes:
   gradedrange,
   isdual,
   nondual
-using NDTensors.LabelledNumbers: LabelledInteger, label, labelled
+using NDTensors.LabelledNumbers: LabelledInteger, label, labelled, labelled_isequal
 using Test: @test, @test_broken, @testset
 struct U1
   n::Int
 end
 GradedAxes.dual(c::U1) = U1(-c.n)
 Base.isless(c1::U1, c2::U1) = c1.n < c2.n
-@testset "dual" begin
-  a = gradedrange([U1(0) => 2, U1(1) => 3])
-  ad = dual(a)
-  @test eltype(ad) == LabelledInteger{Int,U1}
 
-  @test space_isequal(dual(ad), a)
-  @test space_isequal(nondual(ad), a)
-  @test space_isequal(nondual(a), a)
-  @test space_isequal(ad, ad)
-  @test !space_isequal(a, ad)
-  @test !space_isequal(ad, a)
+@testset "AbstractUnitRange" begin
+  a0 = OneToOne()
+  @test !isdual(a0)
+  @test dual(a0) isa OneToOne
+  @test space_isequal(a0, a0)
+  @test labelled_isequal(a0, a0)
+  @test space_isequal(a0, dual(a0))
 
-  @test isdual(ad)
+  a = 1:3
+  ad = dual(a)
   @test !isdual(a)
+  @test !isdual(ad)
+  @test ad isa UnitRange
+  @test space_isequal(ad, a)
 
-  @test blockfirsts(ad) == [labelled(1, U1(0)), labelled(3, U1(-1))]
-  @test blocklasts(ad) == [labelled(2, U1(0)), labelled(5, U1(-1))]
-  @test findblock(ad, 4) == Block(2)
-  @test only(blockaxes(ad)) == Block(1):Block(2)
-  @test blocks(ad) == [labelled(1:2, U1(0)), labelled(3:5, U1(-1))]
-  @test ad[4] == 4
-  @test label(ad[4]) == U1(-1)
-  @test ad[2:4] == 2:4
-  @test ad[2:4] isa UnitRangeDual
-  @test label(ad[2:4][Block(2)]) == U1(-1)
-  @test ad[[2, 4]] == [2, 4]
-  @test label(ad[[2, 4]][2]) == U1(-1)
-  @test ad[Block(2)] == 3:5
-  @test label(ad[Block(2)]) == U1(-1)
-  @test ad[Block(1):Block(2)][Block(2)] == 3:5
-  @test label(ad[Block(1):Block(2)][Block(2)]) == U1(-1)
-  @test ad[[Block(2), Block(1)]][Block(1)] == 3:5
-  @test label(ad[[Block(2), Block(1)]][Block(1)]) == U1(-1)
-  @test ad[[Block(2)[1:2], Block(1)[1:2]]][Block(1)] == 3:4
-  @test label(ad[[Block(2)[1:2], Block(1)[1:2]]][Block(1)]) == U1(-1)
-  @test blocksortperm(a) == [Block(1), Block(2)]
-  @test blocksortperm(ad) == [Block(1), Block(2)]
-  @test blocklength(blockmergesortperm(a)) == 2
-  @test blocklength(blockmergesortperm(ad)) == 2
-  @test blockmergesortperm(a) == [Block(1), Block(2)]
-  @test blockmergesortperm(ad) == [Block(1), Block(2)]
+  a = blockedrange([2, 3])
+  ad = dual(a)
+  @test !isdual(a)
+  @test !isdual(ad)
+  @test ad isa BlockedOneTo
+  @test blockisequal(ad, a)
+end
+
+@testset "GradedUnitRangeDual" begin
+  for a in
+      [gradedrange([U1(0) => 2, U1(1) => 3]), gradedrange([U1(0) => 2, U1(1) => 3])[1:5]]
+    ad = dual(a)
+    @test ad isa GradedUnitRangeDual
+    @test ad isa AbstractGradedUnitRange
+    @test eltype(ad) == LabelledInteger{Int,U1}
+    @test blocklengths(ad) isa Vector
+    @test eltype(blocklengths(ad)) == eltype(blocklengths(a))
+
+    @test space_isequal(dual(ad), a)
+    @test space_isequal(nondual(ad), a)
+    @test space_isequal(nondual(a), a)
+    @test space_isequal(ad, ad)
+    @test !space_isequal(a, ad)
+    @test !space_isequal(ad, a)
+
+    @test isdual(ad)
+    @test !isdual(a)
+    @test axes(Base.Slice(a)) isa Tuple{typeof(a)}
+    @test AbstractUnitRange{Int}(ad) == 1:5
+    b = combine_blockaxes(ad, ad)
+    @test b isa GradedUnitRangeDual
+    @test b == 1:5
+    @test space_isequal(b, ad)
+
+    for x in iterate(ad)
+      @test x == 1
+      @test label(x) == U1(0)
+    end
+    for x in iterate(ad, labelled(3, U1(-1)))
+      @test x == 4
+      @test label(x) == U1(-1)
+    end
+
+    @test blockfirsts(ad) == [labelled(1, U1(0)), labelled(3, U1(-1))]
+    @test blocklasts(ad) == [labelled(2, U1(0)), labelled(5, U1(-1))]
+    @test blocklength(ad) == 2
+    @test blocklengths(ad) == [2, 3]
+    @test blocklabels(ad) == [U1(0), U1(-1)]
+    @test label.(blocklengths(ad)) == [U1(0), U1(-1)]
+    @test findblock(ad, 4) == Block(2)
+    @test only(blockaxes(ad)) == Block(1):Block(2)
+    @test blocks(ad) == [labelled(1:2, U1(0)), labelled(3:5, U1(-1))]
+    @test ad[4] == 4
+    @test label(ad[4]) == U1(-1)
+    @test ad[2:4] == 2:4
+    @test ad[2:4] isa GradedUnitRangeDual
+    @test label(ad[2:4][Block(2)]) == U1(-1)
+    @test ad[[2, 4]] == [2, 4]
+    @test label(ad[[2, 4]][2]) == U1(-1)
+    @test ad[Block(2)] == 3:5
+    @test label(ad[Block(2)]) == U1(-1)
+    @test ad[Block(1):Block(2)][Block(2)] == 3:5
+    @test label(ad[Block(1):Block(2)][Block(2)]) == U1(-1)
+    @test ad[[Block(2), Block(1)]][Block(1)] == 3:5
+    @test label(ad[[Block(2), Block(1)]][Block(1)]) == U1(-1)
+    @test ad[[Block(2)[1:2], Block(1)[1:2]]][Block(1)] == 3:4
+    @test label(ad[[Block(2)[1:2], Block(1)[1:2]]][Block(1)]) == U1(-1)
+    @test blocksortperm(a) == [Block(1), Block(2)]
+    @test blocksortperm(ad) == [Block(1), Block(2)]
+    @test blocklength(blockmergesortperm(a)) == 2
+    @test blocklength(blockmergesortperm(ad)) == 2
+    @test blockmergesortperm(a) == [Block(1), Block(2)]
+    @test blockmergesortperm(ad) == [Block(1), Block(2)]
+
+    @test_broken isdual(ad[Block(1)])
+    @test_broken isdual(ad[Block(1)[1:1]])
+    I = mortar([Block(2)[1:1]])
+    g = ad[I]
+    @test length(g) == 1
+    @test label(first(g)) == U1(-1)
+    @test_broken isdual(g[Block(1)])
+  end
 end
 
 @testset "flip" begin
-  a = gradedrange([U1(0) => 2, U1(1) => 3])
-  ad = dual(a)
-  @test space_isequal(flip(a), dual(gradedrange([U1(0) => 2, U1(-1) => 3])))
-  @test space_isequal(flip(ad), gradedrange([U1(0) => 2, U1(-1) => 3]))
+  for a in
+      [gradedrange([U1(0) => 2, U1(1) => 3]), gradedrange([U1(0) => 2, U1(1) => 3])[1:5]]
+    ad = dual(a)
+    @test space_isequal(flip(a), dual(gradedrange([U1(0) => 2, U1(-1) => 3])))
+    @test space_isequal(flip(ad), gradedrange([U1(0) => 2, U1(-1) => 3]))
 
-  @test blocklabels(a) == [U1(0), U1(1)]
-  @test blocklabels(dual(a)) == [U1(0), U1(-1)]
-  @test blocklabels(flip(a)) == [U1(0), U1(1)]
-  @test blocklabels(flip(dual(a))) == [U1(0), U1(-1)]
-  @test blocklabels(dual(flip(a))) == [U1(0), U1(-1)]
+    @test blocklabels(a) == [U1(0), U1(1)]
+    @test blocklabels(dual(a)) == [U1(0), U1(-1)]
+    @test blocklabels(flip(a)) == [U1(0), U1(1)]
+    @test blocklabels(flip(dual(a))) == [U1(0), U1(-1)]
+    @test blocklabels(dual(flip(a))) == [U1(0), U1(-1)]
 
-  @test blocklengths(a) == [2, 3]
-  @test blocklengths(dual(a)) == [2, 3]
-  @test blocklengths(flip(a)) == [2, 3]
-  @test blocklengths(flip(dual(a))) == [2, 3]
-  @test blocklengths(dual(flip(a))) == [2, 3]
+    @test blocklengths(a) == [2, 3]
+    @test blocklengths(ad) == [2, 3]
+    @test blocklengths(flip(a)) == [2, 3]
+    @test blocklengths(flip(ad)) == [2, 3]
+    @test blocklengths(dual(flip(a))) == [2, 3]
 
-  @test !isdual(a)
-  @test isdual(dual(a))
-  @test isdual(flip(a))
-  @test !isdual(flip(dual(a)))
-  @test !isdual(dual(flip(a)))
+    @test !isdual(a)
+    @test isdual(ad)
+    @test isdual(flip(a))
+    @test !isdual(flip(ad))
+    @test !isdual(dual(flip(a)))
+  end
 end
 end
diff --git a/NDTensors/src/lib/GradedAxes/test/test_tensor_product.jl b/NDTensors/src/lib/GradedAxes/test/test_tensor_product.jl
index 02435b5ba7..99e41454ff 100644
--- a/NDTensors/src/lib/GradedAxes/test/test_tensor_product.jl
+++ b/NDTensors/src/lib/GradedAxes/test/test_tensor_product.jl
@@ -11,11 +11,12 @@ using NDTensors.GradedAxes:
   fusion_product,
   flip,
   gradedrange,
-  labelled_isequal,
   space_isequal,
   isdual,
   tensor_product
 
+using NDTensors.LabelledNumbers: labelled_isequal
+
 struct U1
   n::Int
 end
@@ -27,6 +28,7 @@ GradedAxes.fuse_labels(x::U1, y::U1) = U1(x.n + y.n)
   GradedAxes.fuse_labels(x::String, y::String) = x * y
 
   g0 = OneToOne()
+  @test labelled_isequal(g0, g0)
   @test labelled_isequal(tensor_product(g0, g0), g0)
 
   a = gradedrange(["x" => 2, "y" => 3])
diff --git a/NDTensors/src/lib/LabelledNumbers/src/labelledunitrange.jl b/NDTensors/src/lib/LabelledNumbers/src/labelledunitrange.jl
index a82c666987..03965f62f5 100644
--- a/NDTensors/src/lib/LabelledNumbers/src/labelledunitrange.jl
+++ b/NDTensors/src/lib/LabelledNumbers/src/labelledunitrange.jl
@@ -17,16 +17,6 @@ unlabel_type(::Type{<:LabelledUnitRange{Value}}) where {Value} = Value
 function Base.AbstractUnitRange{T}(a::LabelledUnitRange) where {T}
   return AbstractUnitRange{T}(unlabel(a))
 end
-# Used by `CartesianIndices` constructor.
-# TODO: Seems to only be needed for Julia v1.6, maybe remove once we
-# drop Julia v1.6 support.
-function Base.OrdinalRange{T1,T2}(a::LabelledUnitRange) where {T1,T2<:Integer}
-  return OrdinalRange{T1,T2}(unlabel(a))
-end
-# Fix ambiguity error in Julia v1.10.
-function Base.OrdinalRange{T,T}(a::LabelledUnitRange) where {T<:Integer}
-  return OrdinalRange{T,T}(unlabel(a))
-end
 
 # TODO: Is this a good definition?
 Base.unitrange(a::LabelledUnitRange) = a