How can I understand T(8,128)(2,1) and braces in the error information of padding on TPU?

[mathpluscode]

Hi, I am using JAX with convolutional layers on TPU, and I got the following error message

Program hbm requirement 18.22G:
    global            2.06M
    HLO temp         18.22G (63.3% utilization: Unpadded (11.40G) Padded (18.01G), 1.1% fragmentation (207.55M))

  Largest program allocations in hbm:

  1. Size: 3.00G
     Operator: op_name="jit(<unnamed wrapped function>)/jit(main)/transpose(jvp(unet))/ConvResBlock_15/remat/CheckpointConv_1/conv_general_dilated[window_strides=(1, 1, 1) padding=((1, 1), (1, 1), (1, 1)) lhs_dilat
ion=(1, 1, 1) rhs_dilation=(1, 1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(4, 0, 1, 2, 3), rhs_spec=(4, 0, 1, 2, 3), out_spec=(3, 4, 0, 1, 2)) feature_group_count=1 batch_group_count=1 precision=None pre
ferred_element_type=None]" source_file="/app/ImgX/imgx/model/conv.py" source_line=114
     Shape: bf16[32,128,128,12,64]{0,4,3,2,1:T(8,128)(2,1)}
     Unpadded size: 768.00M
     Extra memory due to padding: 2.25G (4.0x expansion)
     XLA label: copy.621.remat4 = copy(all-gather.80.remat2)
     Allocation type: HLO temp
     ==========================

  2. Size: 768.00M
     Operator: op_name="jit(<unnamed wrapped function>)/jit(main)/transpose(jvp(unet))/ConvResBlock_15/remat/optimization_barrier" source_file="/app/ImgX/imgx/model/conv.py" source_line=114
     Shape: bf16[32,128,128,12,16]{0,4,3,2,1:T(8,128)(2,1)}
     Unpadded size: 192.00M
     Extra memory due to padding: 576.00M (4.0x expansion)
     XLA label: copy.523.remat4 = copy(get-tuple-element.6)
     Allocation type: HLO temp
     ==========================

Here the shape 128,128,12 is the shape of the image and 32 is the batch size. I do understand that on TPU we need some padding

But how can we understand bf16[32,128,128,12,64]{0,4,3,2,1:T(8,128)(2,1)}, particularly what does {0,4,3,2,1} and T(8,128)(2,1) mean here ? Found https://github.com/openxla/xla/blob/main/docs/tiled_layout.md#repeated-tiling,

{0,4,3,2,1} relates to the physical order of dimensions, so [32,128,128,12,64] is actually stored in shape [32,64,12,128,128]. 🤔 I wonder if the tiling of (8,128) corresponds to the axis [32,64] or [64,32]? 🤔 As if it 128 corresponds to 32, that represents the 4x padding, would also explain why it is 4x expansion for both [32,128,128,12,64] and [32,128,128,12,16].

T(8,128)(2,1) means a repeated tiling, first tiled by (8,128) then insider tiled by (2,1), the second tiling should not request padding.

Update: modified my descriptions with some guess from the documents, but would still be happy to get a confirmation! Super thanks for the help in advance!