Repository: tensorflow
Issue: Tensorflow distributed + DTensor approach for large outer product
Labels: stat:awaiting tensorflower,type:support,comp:dist-strat,TF 2.9

### Issue type

Support

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.9

### Custom code

Yes

### OS platform and distribution

MacOS

### Mobile device

N/A

### Python version

3.10

### Bazel version

N/A

### GCC/compiler version

N/A

### CUDA/cuDNN version

N/A

### GPU model and memory

N/A

### Current behavior?

Hey,

I am computing a large matrix of size $(\mathrm{N_{samples}}, \mathrm{N_{parameters}})$ , which corresponds to a jacobian for each sample. I am obtaining this matrix by using a distributed strategy: `tf.distribute.MirroredStrategy` and calculating $N_\mathrm{samples}/N_\mathrm{devices}$ jacobians on all my devices. I then gather along the first axis to construct the full matrix $X$ from which I calculate the matrix $Y = X X^T$ of size $(\mathrm{N_{samples}} ,\mathrm{N_{samples}})$. 

I am looking for a faster way to perform this calculation by sharding along the direction of $N_\mathrm{parameters}$ and performing the calculation like this:

$$
X X^T =\sum^{\mathrm{N_gpus}−1}_{g=0} X_g X_g^T
$$

See this reference for the specific layout:
<img width="887" alt="image" src="https://github.com/user-attachments/assets/47a117c4-906a-4fa3-bf7d-832de810ed7c">

What I tried is to calculate my jacobians with `strategy.run`, gather them and then copy to a mesh to perform the matrix multiplication. However, I feel like this approach is probably not optimal. 

What is the best way to achieve the same within tensorflow's distributed framework? Can I do something similar to MPI's all-to-all primitive to reorder the data across all my devices and perform the matmul described above?

I added a MWE below that outlines my thinking, in my actual code `train_step` involves a complicated calculation and a call to `tf.jacobians` over all my trainable variables.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from tensorflow.experimental import dtensor
import numpy as np

print('TensorFlow version:', tf.__version__)


def configure_virtual_cpus(ncpu):
    phy_devices = tf.config.list_physical_devices('CPU')
    tf.config.set_logical_device_configuration(phy_devices[0], [
        tf.config.LogicalDeviceConfiguration(),
    ] * ncpu)

# Create virtual devices
number_of_devices = 6
configure_virtual_cpus(number_of_devices)
DEVICES = [f'CPU:{i}' for i in range(number_of_devices)]
# Use a mirrorer strategy
strategy = tf.distribute.MirroredStrategy(DEVICES)

number_of_samples_per_device = 4
# Make sure that the number of parameters is divisble by the number of devices
number_of_parameters = 2 * number_of_devices
# Create 1D mesh
mesh = dtensor.create_mesh([("x", number_of_devices), ], devices=DEVICES)
# Shard along Parameter direction
layout = dtensor.Layout([dtensor.UNSHARDED, 'x'], mesh)
tf.random.set_seed(1000)


@tf.function()
def train_step():
    # We create a tensor of (Ns, Np,) as a fictional Jacobian
    return tf.random.uniform((number_of_samples_per_device, number_of_parameters))


@tf.function()
def distributed_step():
    params = strategy.run(train_step)
    # Gathering on axis 1 gives (Ns * N_devices, Np)
    return strategy.gather(params, axis=0)


jacobians = distributed_step()
# Calculate matmul without using a mesh
X_g_full = tf.matmul(jacobians, tf.transpose(jacobians))
# Create a mesh of size (N_devices, )
a = dtensor.copy_to_mesh(jacobians,
                         layout=dtensor.Layout.replicated(layout.mesh, rank=layout.rank))
# Each device gets (Ns * N_devices, Np / N_devices) samples
my_dtensor = dtensor.relayout(a, layout=layout)
for component_tensor in dtensor.unpack(my_dtensor):
    print("Device:", component_tensor.device, ",", component_tensor.shape)
# Perform matmul on mesh.
X_g = tf.matmul(my_dtensor, tf.transpose(my_dtensor))
# They match
print(np.allclose(X_g_full.numpy(), X_g.numpy()))
```


### Relevant log output

```shell
Device: /job:localhost/replica:0/task:0/device:CPU:0 , (24, 2)
Device: /job:localhost/replica:0/task:0/device:CPU:1 , (24, 2)
Device: /job:localhost/replica:0/task:0/device:CPU:2 , (24, 2)
Device: /job:localhost/replica:0/task:0/device:CPU:3 , (24, 2)
Device: /job:localhost/replica:0/task:0/device:CPU:4 , (24, 2)
Device: /job:localhost/replica:0/task:0/device:CPU:5 , (24, 2)
True
```