Add flattenlinear and make it and tryflattenlinear public

Project.toml

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

src/TensorAlgebra.jl

@@ -9,7 +9,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public biperm, bipartition, contractopadd!, data, datatype, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, ungrade, zero!, scale!, permuteddims, PermutedDims"
+            "public biperm, bipartition, contractopadd!, data, datatype, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
         )
     )
 end

src/linearbroadcasted.jl

@@ -399,6 +399,19 @@ function tryflattenlinear(bc::BC.Broadcasted)
     return linearbroadcasted(bc.f, args...)
 end
 
+"""
+    flattenlinear(bc::Broadcasted) -> LinearBroadcasted
+
+Like [`tryflattenlinear`](@ref), but throw an `ArgumentError` when the expression is not
+linear instead of returning `nothing`. The erroring counterpart to `tryflattenlinear`,
+following the `parse`/`tryparse` convention.
+"""
+function flattenlinear(bc)
+    lb = tryflattenlinear(bc)
+    isnothing(lb) && throw(ArgumentError("broadcast expression is not linear"))
+    return lb
+end
+
 # BroadcastStyle for LinearBroadcasted subtypes — delegate to the wrapped array type.
 function BC.BroadcastStyle(::Type{<:ScaledBroadcasted{<:Any, A}}) where {A}
     return BC.BroadcastStyle(A)

test/test_exports.jl

@@ -38,9 +38,10 @@ using Test: @test, @testset
         append!(
             exports,
             [
-                :biperm, :bipartition, :contractopadd!, :data, :datatype, :label_type,
-                :matricizeopperm, :permutedims, :permutedims!, :scalar, :similar_map,
-                :to_range, :tr, :ungrade, :zero!, :scale!, :permuteddims, :PermutedDims,
+                :biperm, :bipartition, :contractopadd!, :data, :datatype, :flattenlinear,
+                :label_type, :matricizeopperm, :permutedims, :permutedims!, :scalar,
+                :similar_map, :to_range, :tr, :tryflattenlinear, :ungrade, :zero!, :scale!,
+                :permuteddims, :PermutedDims,
             ]
         )
     end

test/test_linearbroadcasted.jl

@@ -55,6 +55,21 @@ using Test: @test, @test_throws, @testset
         @test TA.tryflattenlinear(BC.broadcasted(exp, a)) === nothing
         @test TA.tryflattenlinear(BC.broadcasted(+, a, BC.broadcasted(exp, b))) === nothing
     end
+    @testset "flattenlinear" begin
+        a = randn(ComplexF64, 2, 2)
+        b = randn(ComplexF64, 2, 2)
+
+        # Linear expressions convert successfully, matching `tryflattenlinear`
+        @test TA.flattenlinear(BC.broadcasted(*, 2, a)) ≡ linearbroadcasted(*, 2, a)
+        bc = BC.broadcasted(+, BC.broadcasted(*, 2, a), BC.broadcasted(*, 3, b))
+        @test copy(TA.flattenlinear(bc)) ≈ 2a + 3b
+
+        # Nonlinear expression throws instead of returning nothing
+        @test_throws ArgumentError TA.flattenlinear(BC.broadcasted(exp, a))
+        @test_throws ArgumentError TA.flattenlinear(
+            BC.broadcasted(+, a, BC.broadcasted(exp, b))
+        )
+    end
     @testset "linearbroadcasted algebra" begin
         a = randn(ComplexF64, 3, 3)