test BlockedTrivialPermutation

Author: ogauthe

src/blockedpermutation.jl

@@ -83,6 +83,10 @@ function blockedperm(permblocks::Union{Tuple{Vararg{Int}},Int,Ellipsis}...; kwar
   return blockedperm(collect_tuple.(permblocks)...; kwargs...)
 end
 
+function blockedperm(bt::AbstractBlockTuple)
+  return blockedperm(Val(length(bt)), blocks(bt)...)
+end
+
 function _blockedperm_length(::Nothing, specified_perm::Tuple{Vararg{Int}})
   return maximum(specified_perm)
 end
@@ -164,11 +168,19 @@ Base.Tuple(blockedperm::BlockedTrivialPermutation) = flatten_tuples(blocks(block
 
 BlockArrays.blocks(blockedperm::BlockedTrivialPermutation) = getfield(blockedperm, :blocks)
 
+function BlockArrays.blocklengths(
+  ::Type{<:BlockedTrivialPermutation{<:Any,<:Any,Blocks}}
+) where {Blocks}
+  return fieldcount.(fieldtypes(Blocks))
+end
+
+blockedperm(tp::BlockedTrivialPermutation) = tp
+
 function blockedtrivialperm(blocklengths::Tuple{Vararg{Int}})
   return _BlockedTrivialPermutation(blocklengths)
 end
 
-function trivialperm(blockedperm::AbstractBlockedPermutation)
+function trivialperm(blockedperm::AbstractBlockTuple)
   return blockedtrivialperm(blocklengths(blockedperm))
 end
 Base.invperm(blockedperm::BlockedTrivialPermutation) = blockedperm

test/test_basics.jl

@@ -1,6 +1,6 @@
 using EllipsisNotation: var".."
 using LinearAlgebra: norm, qr
-using TensorAlgebra, TensorAlgebra, fusedims, splitdims
+using TensorAlgebra: TensorAlgebra, fusedims, splitdims
 # TODO: Remove dependency on NDTensors, create a GPUTestUtils.jl package.
 using NDTensors: NDTensors
 include(joinpath(pkgdir(NDTensors), "test", "NDTensorsTestUtils", "NDTensorsTestUtils.jl"))

test/test_blockedpermutation.jl

@@ -2,8 +2,11 @@ using Test: @test, @test_broken, @testset
 
 using BlockArrays: blockfirsts, blocklasts, blocklength, blocklengths, blocks
 using Combinatorics: permutations
+using EllipsisNotation: var".."
+using TestExtras: @constinferred
 
-using TensorAlgebra: BlockedTuple, blockedperm, blockedperm_indexin
+using TensorAlgebra:
+  BlockedTrivialPermutation, BlockedTuple, blockedperm, blockedperm_indexin, trivialperm
 
 @testset "BlockedPermutation" begin
   p = blockedperm((3, 4, 5), (2, 1))
@@ -28,7 +31,13 @@ using TensorAlgebra: BlockedTuple, blockedperm, blockedperm_indexin
   @test blockfirsts(p) == (1, 3, 3)
   @test blocklasts(p) == (2, 2, 3)
   @test invperm(p) == blockedperm((3, 2), (), (1,))
-  @test BlockedTuple(p) == BlockedTuple{(2, 0, 1)}((3, 2, 1))
+
+  bt = BlockedTuple{(2, 0, 1)}((3, 2, 1))
+  @test (@constinferred BlockedTuple(p)) == bt
+  @test (@constinferred map(identity, p)) == bt
+  @test (@constinferred p .+ p) == BlockedTuple{(2, 0, 1)}((6, 4, 2))
+  @test (@constinferred blockedperm(p)) == p
+  @test (@constinferred blockedperm(bt)) == p
 
   # Split collection into `BlockedPermutation`.
   p = blockedperm_indexin(("a", "b", "c", "d"), ("c", "a"), ("b", "d"))
@@ -63,3 +72,22 @@ using TensorAlgebra: BlockedTuple, blockedperm, blockedperm_indexin
   p = blockedperm((3, 2), .., 1)
   @test p == blockedperm((3, 2), 1)
 end
+
+@testset "BlockedTrivialPermutation" begin
+  p = blockedperm((3, 2), (), (1,))
+  tp = trivialperm(p)
+
+  @test tp isa BlockedTrivialPermutation
+  @test Tuple(tp) == (1, 2, 3)
+  @test blocklength(tp) == 3
+  @test blocklengths(tp) == (2, 0, 1)
+
+  bt = BlockedTuple{(2, 0, 1)}((1, 2, 3))
+  @test (@constinferred BlockedTuple(tp)) == bt
+  @test (@constinferred blocks(tp)) == blocks(bt)
+  @test (@constinferred map(identity, tp)) == bt
+  @test (@constinferred tp .+ tp) == BlockedTuple{(2, 0, 1)}((2, 4, 6))
+  @test (@constinferred blockedperm(tp)) == tp
+  @test (@constinferred trivialperm(tp)) == tp
+  @test (@constinferred trivialperm(bt)) == tp
+end