generic broadcast

Author: ogauthe

Project.toml

@@ -9,6 +9,7 @@ BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
@@ -23,4 +24,5 @@ EllipsisNotation = "1.8.0"
 GradedUnitRanges = "0.1.0"
 LinearAlgebra = "1.10"
 TupleTools = "1.6.0"
+TypeParameterAccessors = "0.2.1"
 julia = "1.10"

src/blockedtuple.jl

@@ -3,6 +3,8 @@
 
 using BlockArrays: Block, BlockArrays, BlockIndexRange, BlockRange, blockedrange
 
+using TypeParameterAccessors: unspecify_type_parameters
+
 #
 # ==================================  AbstractBlockTuple  ==================================
 #
@@ -18,7 +20,12 @@ Base.firstindex(::AbstractBlockTuple) = 1
 Base.getindex(bt::AbstractBlockTuple, i::Integer) = Tuple(bt)[i]
 Base.getindex(bt::AbstractBlockTuple, r::AbstractUnitRange) = Tuple(bt)[r]
 Base.getindex(bt::AbstractBlockTuple, b::Block{1}) = blocks(bt)[Int(b)]
-Base.getindex(bt::AbstractBlockTuple, br::BlockRange{1}) = blocks(bt)[Int.(br)]
+function Base.getindex(bt::AbstractBlockTuple, br::BlockRange{1})
+  r = Int.(br)
+  T = unspecify_type_parameters(typeof(bt))
+  flat = Tuple(bt)[blockfirsts(bt)[first(r)]:blocklasts(bt)[last(r)]]
+  return T{blocklengths(bt)[r]}(flat)
+end
 function Base.getindex(bt::AbstractBlockTuple, bi::BlockIndexRange{1})
   return bt[Block(bi)][only(bi.indices)]
 end
@@ -32,19 +39,21 @@ Base.lastindex(bt::AbstractBlockTuple) = length(bt)
 
 # Broadcast interface
 Base.broadcastable(bt::AbstractBlockTuple) = bt
-struct BlockedTupleBroadcastStyle{BlockLengths} <: Broadcast.BroadcastStyle end
-function Base.BroadcastStyle(type::Type{<:AbstractBlockTuple})
-  return BlockedTupleBroadcastStyle{blocklengths(type)}()
+struct AbstractBlockTupleBroadcastStyle{BlockLengths,BT} <: Broadcast.BroadcastStyle end
+function Base.BroadcastStyle(T::Type{<:AbstractBlockTuple})
+  return AbstractBlockTupleBroadcastStyle{blocklengths(T),unspecify_type_parameters(T)}()
 end
 
 # BroadcastStyle is not called for two identical styles
-function Base.BroadcastStyle(::BlockedTupleBroadcastStyle, ::BlockedTupleBroadcastStyle)
+function Base.BroadcastStyle(
+  ::AbstractBlockTupleBroadcastStyle, ::AbstractBlockTupleBroadcastStyle
+)
   throw(DimensionMismatch("Incompatible blocks"))
 end
 function Base.copy(
-  bc::Broadcast.Broadcasted{BlockedTupleBroadcastStyle{BlockLengths}}
-) where {BlockLengths}
-  return BlockedTuple{BlockLengths}(bc.f.((Tuple.(bc.args))...))
+  bc::Broadcast.Broadcasted{AbstractBlockTupleBroadcastStyle{BlockLengths,BT}}
+) where {BlockLengths,BT}
+  return BT{BlockLengths}(bc.f.((Tuple.(bc.args))...))
 end
 
 # BlockArrays interface

test/test_blockedtuple.jl

@@ -25,7 +25,7 @@ using TensorAlgebra: BlockedTuple, tuplemortar
   # it is hard to make bt[Block(1)] type stable as compile-time knowledge of 1 is lost in Block
   @test bt[Block(1)] == blocks(bt)[1]
   @test bt[Block(2)] == blocks(bt)[2]
-  @test bt[Block(1):Block(2)] == blocks(bt)[1:2]
+  @test bt[Block(1):Block(2)] == tuplemortar((true,), ('a', 2))
   @test bt[Block(2)[1:2]] == ('a', 2)
   @test bt[2:4] == ('a', 2, "b")