From d132d0bb7ef58ac6aea90b4b294b1034b420942d Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 13 Mar 2025 10:30:59 -0400
Subject: [PATCH 01/11] Define random_unitary constructor

---
 Project.toml                                  |  8 ++-
 .../TensorAlgebraGradedUnitRangesExt.jl       | 25 ++++++-
 src/TensorAlgebra.jl                          |  1 +
 src/random_unitary.jl                         | 69 +++++++++++++++++++
 4 files changed, 98 insertions(+), 5 deletions(-)
 create mode 100644 src/random_unitary.jl

diff --git a/Project.toml b/Project.toml
index 99aa433..f997b51 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,13 +1,14 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
@@ -22,7 +23,8 @@ ArrayLayouts = "1.10.4"
 BlockArrays = "1.2.0"
 EllipsisNotation = "1.8.0"
 GradedUnitRanges = "0.1.0"
-LinearAlgebra = "1.10"
+LinearAlgebra = "1.10.0"
+Random = "1.10.0"
 TupleTools = "1.6.0"
-TypeParameterAccessors = "0.2.1, 0.3"
+TypeParameterAccessors = "0.2.1, 0.3.0"
 julia = "1.10"
diff --git a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
index fe7caa2..f38c722 100644
--- a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
@@ -1,8 +1,29 @@
 module TensorAlgebraGradedUnitRangesExt
-using GradedUnitRanges: AbstractGradedUnitRange, tensor_product
-using TensorAlgebra: TensorAlgebra
+
+using GradedUnitRanges: AbstractGradedUnitRange, dual, tensor_product
+using GradedUnitRanges.BlockArrays: Block, blocklengths, blocksize
+using Random: AbstractRNG
+using TensorAlgebra: TensorAlgebra, random_unitary
 
 function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
   return tensor_product(a1, a2)
 end
+
+function TensorAlgebra.dual(a::AbstractGradedUnitRange)
+  return dual(a)
+end
+
+function TensorAlgebra.random_unitary(
+  rng::AbstractRNG,
+  elt::Type,
+  ax::Tuple{AbstractGradedUnitRange},
+)
+  a = zeros(elt, dual.(ax)..., ax...)
+  # TODO: Define `blockdiagindices`.
+  for i in 1:minimum(blocksize(a))
+    a[Block(i, i)] = random_unitary(rng, elt, Int(blocklengths(only(ax))[i]))
+  end
+  return a
+end
+
 end
diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl
index caa2cc5..87f847c 100644
--- a/src/TensorAlgebra.jl
+++ b/src/TensorAlgebra.jl
@@ -13,5 +13,6 @@ include("contract/blockedperms.jl")
 include("contract/allocate_output.jl")
 include("contract/contract_matricize/contract.jl")
 include("factorizations.jl")
+include("random_unitary.jl")
 
 end
diff --git a/src/random_unitary.jl b/src/random_unitary.jl
new file mode 100644
index 0000000..21a9196
--- /dev/null
+++ b/src/random_unitary.jl
@@ -0,0 +1,69 @@
+# Version of `sign` that returns one
+# if `x == 0`.
+function nonzero_sign(x)
+  iszero(x) && return one(x)
+  return sign(x)
+end
+
+using LinearAlgebra: LinearAlgebra, Diagonal, diag
+function qr_positive(M::AbstractMatrix)
+  Q, R = LinearAlgebra.qr(M)
+  Q′ = typeof(R)(Q)
+  signs = nonzero_sign.(diag(R))
+  Q′ = Q′ * Diagonal(signs)
+  R = Diagonal(conj.(signs)) * R
+  return Q′, R
+end
+
+using Random: Random, AbstractRNG
+
+dual(x) = x
+
+function random_unitary(
+  rng::AbstractRNG,
+  elt::Type,
+  ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}},
+)
+  ax_fused = ⊗(ax...)
+  a_fused = random_unitary(rng, elt, ax_fused)
+  return splitdims(a_fused, dual.(ax), ax)
+end
+
+# Copy of `Base.to_dim`:
+# https://github.com/JuliaLang/julia/blob/1431bec1bcd205f181ca2a3f1c314247b64076df/base/array.jl#L439-L440
+to_dim(d::Integer) = d
+to_dim(d::Base.OneTo) = last(d)
+
+# Matrix version.
+function random_unitary(
+  rng::AbstractRNG,
+  elt::Type,
+  ax::Tuple{AbstractUnitRange},
+)
+  return random_unitary(rng, elt, map(to_dim, ax))
+end
+
+function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Integer})
+  Q, _ = qr_positive(randn(rng, elt, (dims..., dims...)))
+  return Q
+end
+
+# Canonicalizing other kinds of inputs.
+function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}})
+  return random_unitary(Random.default_rng(), elt, map(to_axis, dims))
+end
+function random_unitary(elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}})
+  return random_unitary(Random.default_rng(), elt, dims)
+end
+function random_unitary(rng::AbstractRNG, elt::Type, dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(rng, elt, dims)
+end
+function random_unitary(elt::Type, dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(Random.default_rng(), elt, dims)
+end
+function random_unitary(rng::AbstractRNG, dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(rng, Float64, dims)
+end
+function random_unitary(dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(Random.default_rng(), Float64, dims)
+end

From 8a1b555ac7f7b7de8923488b500e2a59a6aef22b Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 09:54:38 -0400
Subject: [PATCH 02/11] Use MatrixAlgebraKit

---
 Project.toml                                  |  2 +
 .../TensorAlgebraGradedUnitRangesExt.jl       | 14 +++----
 src/fusedims.jl                               |  2 +
 src/random_unitary.jl                         | 40 +++++--------------
 4 files changed, 21 insertions(+), 37 deletions(-)

diff --git a/Project.toml b/Project.toml
index f997b51..47b5a5e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -8,6 +8,7 @@ ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
@@ -24,6 +25,7 @@ BlockArrays = "1.2.0"
 EllipsisNotation = "1.8.0"
 GradedUnitRanges = "0.1.0"
 LinearAlgebra = "1.10.0"
+MatrixAlgebraKit = "0.1.1"
 Random = "1.10.0"
 TupleTools = "1.6.0"
 TypeParameterAccessors = "0.2.1, 0.3.0"
diff --git a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
index f38c722..cd243bf 100644
--- a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
@@ -14,14 +14,14 @@ function TensorAlgebra.dual(a::AbstractGradedUnitRange)
 end
 
 function TensorAlgebra.random_unitary(
-  rng::AbstractRNG,
-  elt::Type,
-  ax::Tuple{AbstractGradedUnitRange},
+  rng::AbstractRNG, elt::Type, axs::Tuple{AbstractGradedUnitRange}
 )
-  a = zeros(elt, dual.(ax)..., ax...)
-  # TODO: Define `blockdiagindices`.
-  for i in 1:minimum(blocksize(a))
-    a[Block(i, i)] = random_unitary(rng, elt, Int(blocklengths(only(ax))[i]))
+  ax = only(axs)
+  a = zeros(elt, dual(ax), ax)
+  # TODO: Define and use `blockdiagindices`
+  # or `blockdiaglength`.
+  for i in 1:blocksize(a, 1)
+    a[Block(i, i)] = random_unitary(rng, elt, Int(blocklengths(ax)[i]))
   end
   return a
 end
diff --git a/src/fusedims.jl b/src/fusedims.jl
index 38205c8..2dd7d1e 100644
--- a/src/fusedims.jl
+++ b/src/fusedims.jl
@@ -21,6 +21,8 @@ function fusedims(::ReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...)
   return reshape(a, axes)
 end
 
+dual(x) = x
+
 ⊗(a::AbstractUnitRange) = a
 function ⊗(a1::AbstractUnitRange, a2::AbstractUnitRange, as::AbstractUnitRange...)
   return ⊗(a1, ⊗(a2, as...))
diff --git a/src/random_unitary.jl b/src/random_unitary.jl
index 21a9196..a7e7e40 100644
--- a/src/random_unitary.jl
+++ b/src/random_unitary.jl
@@ -1,28 +1,7 @@
-# Version of `sign` that returns one
-# if `x == 0`.
-function nonzero_sign(x)
-  iszero(x) && return one(x)
-  return sign(x)
-end
-
-using LinearAlgebra: LinearAlgebra, Diagonal, diag
-function qr_positive(M::AbstractMatrix)
-  Q, R = LinearAlgebra.qr(M)
-  Q′ = typeof(R)(Q)
-  signs = nonzero_sign.(diag(R))
-  Q′ = Q′ * Diagonal(signs)
-  R = Diagonal(conj.(signs)) * R
-  return Q′, R
-end
-
 using Random: Random, AbstractRNG
 
-dual(x) = x
-
 function random_unitary(
-  rng::AbstractRNG,
-  elt::Type,
-  ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}},
+  rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}
 )
   ax_fused = ⊗(ax...)
   a_fused = random_unitary(rng, elt, ax_fused)
@@ -35,27 +14,28 @@ to_dim(d::Integer) = d
 to_dim(d::Base.OneTo) = last(d)
 
 # Matrix version.
-function random_unitary(
-  rng::AbstractRNG,
-  elt::Type,
-  ax::Tuple{AbstractUnitRange},
-)
+function random_unitary(rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange})
   return random_unitary(rng, elt, map(to_dim, ax))
 end
 
+using MatrixAlgebraKit: qr_full!
 function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Integer})
-  Q, _ = qr_positive(randn(rng, elt, (dims..., dims...)))
+  Q, _ = qr_full!(randn(rng, elt, (dims..., dims...)); positive=true)
   return Q
 end
 
 # Canonicalizing other kinds of inputs.
-function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}})
+function random_unitary(
+  rng::AbstractRNG, elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}}
+)
   return random_unitary(Random.default_rng(), elt, map(to_axis, dims))
 end
 function random_unitary(elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}})
   return random_unitary(Random.default_rng(), elt, dims)
 end
-function random_unitary(rng::AbstractRNG, elt::Type, dims::Union{AbstractUnitRange,Integer}...)
+function random_unitary(
+  rng::AbstractRNG, elt::Type, dims::Union{AbstractUnitRange,Integer}...
+)
   return random_unitary(rng, elt, dims)
 end
 function random_unitary(elt::Type, dims::Union{AbstractUnitRange,Integer}...)

From 6a849a7752f1fbc68ada04a70e12faaf6cfa26e6 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 14:48:52 -0400
Subject: [PATCH 03/11] Refactor in terms of in-place random_unitary

---
 Project.toml                                  |  4 +-
 ...braBlockSparseArraysGradedUnitRangesExt.jl | 33 ++++++++++++++++
 .../TensorAlgebraGradedUnitRangesExt.jl       | 29 --------------
 src/fusedims.jl                               |  2 -
 src/random_unitary.jl                         | 39 ++++++++++---------
 test/Project.toml                             |  1 +
 6 files changed, 58 insertions(+), 50 deletions(-)
 create mode 100644 ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
 delete mode 100644 ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl

diff --git a/Project.toml b/Project.toml
index 47b5a5e..eea9800 100644
--- a/Project.toml
+++ b/Project.toml
@@ -14,14 +14,16 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
+BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
 
 [extensions]
-TensorAlgebraGradedUnitRangesExt = "GradedUnitRanges"
+TensorAlgebraBlockSparseArraysGradedUnitRangesExt = ["BlockSparseArrays", "GradedUnitRanges"]
 
 [compat]
 ArrayLayouts = "1.10.4"
 BlockArrays = "1.2.0"
+BlockSparseArrays = "0.3.6"
 EllipsisNotation = "1.8.0"
 GradedUnitRanges = "0.1.0"
 LinearAlgebra = "1.10.0"
diff --git a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
new file mode 100644
index 0000000..ff4e8c3
--- /dev/null
+++ b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
@@ -0,0 +1,33 @@
+module TensorAlgebraBlockSparseArraysGradedUnitRangesExt
+
+using BlockArrays: Block, blocksize
+using BlockSparseArrays: BlockSparseMatrix, @view!
+using GradedUnitRanges: AbstractGradedUnitRange, dual, tensor_product
+using Random: AbstractRNG
+using TensorAlgebra: TensorAlgebra, random_unitary!
+
+function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
+  return tensor_product(a1, a2)
+end
+
+function TensorAlgebra.square_zero_map(
+  elt::Type, ax::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}
+)
+  return BlockSparseArray{elt}(undef, (dual.(ax)..., ax...))
+end
+
+function TensorAlgebra.random_unitary!(
+  rng::AbstractRNG,
+  a::BlockSparseMatrix{
+    <:Any,<:Any,<:Any,<:Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}
+  },
+)
+  # TODO: Define and use `blockdiagindices`
+  # or `blockdiaglength`.
+  for i in 1:blocksize(a, 1)
+    a[Block(i, i)] = random_unitary!(rng, @view!(a[Block(i, i)]))
+  end
+  return a
+end
+
+end
diff --git a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
deleted file mode 100644
index cd243bf..0000000
--- a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
+++ /dev/null
@@ -1,29 +0,0 @@
-module TensorAlgebraGradedUnitRangesExt
-
-using GradedUnitRanges: AbstractGradedUnitRange, dual, tensor_product
-using GradedUnitRanges.BlockArrays: Block, blocklengths, blocksize
-using Random: AbstractRNG
-using TensorAlgebra: TensorAlgebra, random_unitary
-
-function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
-  return tensor_product(a1, a2)
-end
-
-function TensorAlgebra.dual(a::AbstractGradedUnitRange)
-  return dual(a)
-end
-
-function TensorAlgebra.random_unitary(
-  rng::AbstractRNG, elt::Type, axs::Tuple{AbstractGradedUnitRange}
-)
-  ax = only(axs)
-  a = zeros(elt, dual(ax), ax)
-  # TODO: Define and use `blockdiagindices`
-  # or `blockdiaglength`.
-  for i in 1:blocksize(a, 1)
-    a[Block(i, i)] = random_unitary(rng, elt, Int(blocklengths(ax)[i]))
-  end
-  return a
-end
-
-end
diff --git a/src/fusedims.jl b/src/fusedims.jl
index 2dd7d1e..38205c8 100644
--- a/src/fusedims.jl
+++ b/src/fusedims.jl
@@ -21,8 +21,6 @@ function fusedims(::ReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...)
   return reshape(a, axes)
 end
 
-dual(x) = x
-
 ⊗(a::AbstractUnitRange) = a
 function ⊗(a1::AbstractUnitRange, a2::AbstractUnitRange, as::AbstractUnitRange...)
   return ⊗(a1, ⊗(a2, as...))
diff --git a/src/random_unitary.jl b/src/random_unitary.jl
index a7e7e40..83523a3 100644
--- a/src/random_unitary.jl
+++ b/src/random_unitary.jl
@@ -1,29 +1,32 @@
-using Random: Random, AbstractRNG
+using MatrixAlgebraKit: qr_full!
+using Random: Random, AbstractRNG, randn!
 
-function random_unitary(
-  rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}
-)
-  ax_fused = ⊗(ax...)
-  a_fused = random_unitary(rng, elt, ax_fused)
-  return splitdims(a_fused, dual.(ax), ax)
+function square_zero_map(elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}})
+  return zeros(elt, (ax..., ax...))
 end
 
-# Copy of `Base.to_dim`:
-# https://github.com/JuliaLang/julia/blob/1431bec1bcd205f181ca2a3f1c314247b64076df/base/array.jl#L439-L440
-to_dim(d::Integer) = d
-to_dim(d::Base.OneTo) = last(d)
-
-# Matrix version.
-function random_unitary(rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange})
-  return random_unitary(rng, elt, map(to_dim, ax))
+using EllipsisNotation: : .. function random_unitary!(rng::AbstractRNG, a::AbstractArray)
+  @assert iseven(ndims(a))
+  ndims_codomain = ndims(a) ÷ 2
+  biperm = blockedperm(ntuple(identity, ndims(a)), (ndims_codomain, ndims_codomain))
+  a_mat = fusedims(a, biperm)
+  r_mat = random_unitary!(rng, a_mat)
+  splitdims!(a, r_mat, biperm)
+  return a
 end
 
-using MatrixAlgebraKit: qr_full!
-function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Integer})
-  Q, _ = qr_full!(randn(rng, elt, (dims..., dims...)); positive=true)
+function random_unitary!(rng::AbstractRNG, a::AbstractMatrix)
+  a_r = randn!(rng, a)
+  Q, _ = qr_full!(randn!(rng, a); positive=true)
   return Q
 end
 
+function random_unitary(
+  rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}
+)
+  return random_unitary!(rng, square_zero_map(elt, ax))
+end
+
 # Canonicalizing other kinds of inputs.
 function random_unitary(
   rng::AbstractRNG, elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}}
diff --git a/test/Project.toml b/test/Project.toml
index 7258ca4..ce0519f 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,6 +1,7 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"

From 219ea49f4fa7c85c46e2a12b6abdb0c63d7bfb10 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 14:51:41 -0400
Subject: [PATCH 04/11] Fix import

---
 src/random_unitary.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/random_unitary.jl b/src/random_unitary.jl
index 83523a3..e398d43 100644
--- a/src/random_unitary.jl
+++ b/src/random_unitary.jl
@@ -5,7 +5,10 @@ function square_zero_map(elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractU
   return zeros(elt, (ax..., ax...))
 end
 
-using EllipsisNotation: : .. function random_unitary!(rng::AbstractRNG, a::AbstractArray)
+# Imports `..` into the namespace.
+using EllipsisNotation
+
+function random_unitary!(rng::AbstractRNG, a::AbstractArray)
   @assert iseven(ndims(a))
   ndims_codomain = ndims(a) ÷ 2
   biperm = blockedperm(ntuple(identity, ndims(a)), (ndims_codomain, ndims_codomain))

From 2e0e67288c2b0403d96ebc13323c5585beadec36 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 14:59:06 -0400
Subject: [PATCH 05/11] Simplify

---
 .../TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl  | 2 +-
 src/random_unitary.jl                                     | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
index ff4e8c3..b385261 100644
--- a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
@@ -25,7 +25,7 @@ function TensorAlgebra.random_unitary!(
   # TODO: Define and use `blockdiagindices`
   # or `blockdiaglength`.
   for i in 1:blocksize(a, 1)
-    a[Block(i, i)] = random_unitary!(rng, @view!(a[Block(i, i)]))
+    random_unitary!(rng, @view!(a[Block(i, i)]))
   end
   return a
 end
diff --git a/src/random_unitary.jl b/src/random_unitary.jl
index e398d43..9c04275 100644
--- a/src/random_unitary.jl
+++ b/src/random_unitary.jl
@@ -13,15 +13,15 @@ function random_unitary!(rng::AbstractRNG, a::AbstractArray)
   ndims_codomain = ndims(a) ÷ 2
   biperm = blockedperm(ntuple(identity, ndims(a)), (ndims_codomain, ndims_codomain))
   a_mat = fusedims(a, biperm)
-  r_mat = random_unitary!(rng, a_mat)
-  splitdims!(a, r_mat, biperm)
+  random_unitary!(rng, a_mat)
+  splitdims!(a, a_mat, biperm)
   return a
 end
 
 function random_unitary!(rng::AbstractRNG, a::AbstractMatrix)
-  a_r = randn!(rng, a)
   Q, _ = qr_full!(randn!(rng, a); positive=true)
-  return Q
+  a .= Q
+  return a
 end
 
 function random_unitary(

From e47a54499c5bdeee73821540e8dd80d1a67c84e4 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 15:03:24 -0400
Subject: [PATCH 06/11] Reorganize extensions

---
 Project.toml                                           |  1 +
 .../TensorAlgebraGradedUnitRangesExt.jl                | 10 ++++++++++
 2 files changed, 11 insertions(+)
 create mode 100644 ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl

diff --git a/Project.toml b/Project.toml
index eea9800..66024f4 100644
--- a/Project.toml
+++ b/Project.toml
@@ -19,6 +19,7 @@ GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
 
 [extensions]
 TensorAlgebraBlockSparseArraysGradedUnitRangesExt = ["BlockSparseArrays", "GradedUnitRanges"]
+TensorAlgebraGradedUnitRangesExt = "GradedUnitRanges"
 
 [compat]
 ArrayLayouts = "1.10.4"
diff --git a/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
new file mode 100644
index 0000000..c9d2ad5
--- /dev/null
+++ b/ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl
@@ -0,0 +1,10 @@
+module TensorAlgebraGradedUnitRangesExt
+
+using GradedUnitRanges: AbstractGradedUnitRange, tensor_product
+using TensorAlgebra: TensorAlgebra
+
+function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
+  return tensor_product(a1, a2)
+end
+
+end

From bab22acfc7434c403aa7cd68d53139bebc483bc2 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 15:08:28 -0400
Subject: [PATCH 07/11] Delete code duplication

---
 .../TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl  | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
index b385261..9696771 100644
--- a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
@@ -2,14 +2,10 @@ module TensorAlgebraBlockSparseArraysGradedUnitRangesExt
 
 using BlockArrays: Block, blocksize
 using BlockSparseArrays: BlockSparseMatrix, @view!
-using GradedUnitRanges: AbstractGradedUnitRange, dual, tensor_product
+using GradedUnitRanges: AbstractGradedUnitRange, dual
 using Random: AbstractRNG
 using TensorAlgebra: TensorAlgebra, random_unitary!
 
-function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
-  return tensor_product(a1, a2)
-end
-
 function TensorAlgebra.square_zero_map(
   elt::Type, ax::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}
 )
@@ -19,7 +15,7 @@ end
 function TensorAlgebra.random_unitary!(
   rng::AbstractRNG,
   a::BlockSparseMatrix{
-    <:Any,<:Any,<:Any,<:Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}
+    <:Any,<:Any,<:Any,<:NTuple{2,AbstractGradedUnitRange}
   },
 )
   # TODO: Define and use `blockdiagindices`

From eb86d464b689a9968f7d623b96038582310cc4d1 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 15:14:09 -0400
Subject: [PATCH 08/11] Format

---
 .../TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl      | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
index 9696771..0c00935 100644
--- a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
@@ -14,9 +14,7 @@ end
 
 function TensorAlgebra.random_unitary!(
   rng::AbstractRNG,
-  a::BlockSparseMatrix{
-    <:Any,<:Any,<:Any,<:NTuple{2,AbstractGradedUnitRange}
-  },
+  a::BlockSparseMatrix{<:Any,<:Any,<:Any,<:NTuple{2,AbstractGradedUnitRange}},
 )
   # TODO: Define and use `blockdiagindices`
   # or `blockdiaglength`.

From 067878503f168afc1232da33c66fb22dc2a6ac4d Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 17:20:54 -0400
Subject: [PATCH 09/11] Check square

---
 .../TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl     | 5 +++--
 src/random_unitary.jl                                        | 3 ---
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
index 0c00935..fb05ff5 100644
--- a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
@@ -1,8 +1,8 @@
 module TensorAlgebraBlockSparseArraysGradedUnitRangesExt
 
 using BlockArrays: Block, blocksize
-using BlockSparseArrays: BlockSparseMatrix, @view!
-using GradedUnitRanges: AbstractGradedUnitRange, dual
+using BlockSparseArrays: BlockSparseArray, BlockSparseMatrix, @view!
+using GradedUnitRanges: AbstractGradedUnitRange, dual, space_isequal
 using Random: AbstractRNG
 using TensorAlgebra: TensorAlgebra, random_unitary!
 
@@ -16,6 +16,7 @@ function TensorAlgebra.random_unitary!(
   rng::AbstractRNG,
   a::BlockSparseMatrix{<:Any,<:Any,<:Any,<:NTuple{2,AbstractGradedUnitRange}},
 )
+  space_isequal(axes(a, 1), dual(axes(a, 2))) || throw(ArgumentError("Codomain and domain spaces must be equal."))
   # TODO: Define and use `blockdiagindices`
   # or `blockdiaglength`.
   for i in 1:blocksize(a, 1)
diff --git a/src/random_unitary.jl b/src/random_unitary.jl
index 9c04275..ab00c28 100644
--- a/src/random_unitary.jl
+++ b/src/random_unitary.jl
@@ -5,9 +5,6 @@ function square_zero_map(elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractU
   return zeros(elt, (ax..., ax...))
 end
 
-# Imports `..` into the namespace.
-using EllipsisNotation
-
 function random_unitary!(rng::AbstractRNG, a::AbstractArray)
   @assert iseven(ndims(a))
   ndims_codomain = ndims(a) ÷ 2

From 9440534bffc5a835da5b326d6fb337ef11ac5b01 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 14 Mar 2025 17:23:36 -0400
Subject: [PATCH 10/11] Format

---
 .../TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl       | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
index fb05ff5..12044bd 100644
--- a/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
+++ b/ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl
@@ -16,7 +16,8 @@ function TensorAlgebra.random_unitary!(
   rng::AbstractRNG,
   a::BlockSparseMatrix{<:Any,<:Any,<:Any,<:NTuple{2,AbstractGradedUnitRange}},
 )
-  space_isequal(axes(a, 1), dual(axes(a, 2))) || throw(ArgumentError("Codomain and domain spaces must be equal."))
+  space_isequal(axes(a, 1), dual(axes(a, 2))) ||
+    throw(ArgumentError("Codomain and domain spaces must be equal."))
   # TODO: Define and use `blockdiagindices`
   # or `blockdiaglength`.
   for i in 1:blocksize(a, 1)

From 879ef7db5414c0d4dbab4cbb66f2037c26db9ea5 Mon Sep 17 00:00:00 2001
From: Matt Fishman <mtfishman@users.noreply.github.com>
Date: Thu, 20 Mar 2025 17:44:22 -0400
Subject: [PATCH 11/11] Properly pass along RNG

---
 src/random_unitary.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/random_unitary.jl b/src/random_unitary.jl
index ab00c28..7d57f79 100644
--- a/src/random_unitary.jl
+++ b/src/random_unitary.jl
@@ -31,7 +31,7 @@ end
 function random_unitary(
   rng::AbstractRNG, elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}}
 )
-  return random_unitary(Random.default_rng(), elt, map(to_axis, dims))
+  return random_unitary(rng, elt, map(to_axis, dims))
 end
 function random_unitary(elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}})
   return random_unitary(Random.default_rng(), elt, dims)