[mosaic TPU][apply vector layout for matmul ] why row dimension value of matmul's lhs operand shape is large, but not restrict to mxc_noncontracting_size?

[jackiechen19940708]

Hi, I have a question about matmul apply vector layout rule. when applying layout, I see that lhs_cols (https://github.com/jax-ml/jax/blob/main/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc#L2721 ) use target_lhs_vregs.dim(0) as the row dimension value, this can be larger (for example 1024 @ 32bit) than mxc_noncontracting_size(for example 128 @ 32bit), so why lhs do not tile to 128,128 @ 32bit while rhs does so.

is google tpu hardware supporting such format (lhs > 128 @32bit) matmul instruction? or it will be tiled in some following passes which is not open sourced?