What is the default multicore behaviour on TPUs?

[oluwatimilehin]

Hello, I'm benchmarking some simple kernels like jax.lax.add, jax.lax.subtract, etc., on the TPUv4, and it looks like the kernel gets executed on both cores:

The above is the trace from running jax.lax.abs on a 1x4096x256 tensor. It looks like both cores execute the kernel at the same time.

What I find confusing, though, is that the tensor does not seem to be split between the tensor cores. For example, selecting one of the blocks gives the same HLO for both cores:

%abs.2 = s32[1,4096,256]{2,1,0:T(8,128)} abs(s32[1,4096,256]{2,1,0:T(8,128)} %Arg_0.1) 

I'm not using any Jax primitives to split the work, and I don't understand why it's running the operation on both cores. From the profile, it also looks like CMEM is not being used at all, as the read and write bandwidth utilization for both is zero. Could you please explain the behaviour?