def unbatch(batched_values):
    """Unbatch values across leading axis and return a list of elements."""
    b1, b2 = batched_values.shape[0], batched_values.shape[1]
    results = []
    for v_array in jnp.split(batched_values, indices_or_sections=b1, axis=0):
        v_array = jnp.squeeze(v_array)
        # b2 = batches (number of preconditioner computation) per core.
        if b2 > 1:
            for v in jnp.split(v_array, indices_or_sections=b2, axis=0):
                results.append(jnp.squeeze(v))
        else:
            results.append(v_array)
    return results
Unbatch values across leading axis and return a list of elements.
tools/train/scalable_shampoo/distributed_shampoo.py