Symmetric mesh, but output sharded differently

[mattiasmar]

I extracted the test test_pjit_gda_mixed_inputs from pjit_test.py unit tests and slightly edited the mesh to be symmetric.
Could you explain why in this setup the two output values (out1 & out2) get sharded differently?

out1:
[[ 1  3  5  7]
 [ 3 13 23 33]
 [ 5 23 41 59]
 [ 7 33 59 85]]
[[  9  11  13  15]
 [ 43  53  63  73]
 [ 77  95 113 131]
 [111 137 163 189]]
[[  9  43  77 111]
 [ 11  53  95 137]
 [ 13  63 113 163]
 [ 15  73 131 189]]
[[145 179 213 247]
 [179 221 263 305]
 [213 263 313 363]
 [247 305 363 421]]
out2:
[[ 1  3  5  7  9 11 13 15]
 [ 3 13 23 33 43 53 63 73]]
[[  5  23  41  59  77  95 113 131]
 [  7  33  59  85 111 137 163 189]]
[[  9  43  77 111 145 179 213 247]
 [ 11  53  95 137 179 221 263 305]]
[[ 13  63 113 163 213 263 313 363]
 [ 15  73 131 189 247 305 363 421]]

global_mesh = create_global_mesh((2, 2), ('x', 'y'))
global_input_shape = (8, 2)
mesh_axes = P('x', 'y')
input_data = np.arange(
    prod(global_input_shape)).reshape(global_input_shape)

def cb(index):
    return input_data[index]

gda_obj = global_device_array.GlobalDeviceArray.from_callback(
    global_input_shape, global_mesh, mesh_axes, cb)


with jax._src.config.parallel_functions_output_gda(True):
    @partial(pjit,
             in_axis_resources=(P('x', 'y'), P('x', 'y')),
             out_axis_resources=(P('x', 'y'), P(('x', 'y'))))
    def f(x, y):
        return x @ x.T, y @ y.T

    expected_matrix_mul = input_data @ input_data.T
    with jtu.with_mesh([('x', 2), ('y', 2)]):
        out1, out2 = f(gda_obj, input_data)