Implemented SAMPLED_ADDMM
You might be able to do this with existing ops:
- `tf.gather_nd` to extract the appropriate rows/columns of `mat_a` and `mat_b`.
- `tf.matmul` to do the dot-product
- scale and sum to the sparse input values

The benefit of doing it this way is XLA should just work out of the box, and you wouldn't need to support a new op. @cantonios Thanks for the feedback! Will look more into it over the next couple days. Working on a shape function now. > You might be able to do this with existing ops:
> 
> * `tf.gather_nd` to extract the appropriate rows/columns of `mat_a` and `mat_b`.
> * `tf.matmul` to do the dot-product
> * scale and sum to the sparse input values
> 
> The benefit of doing it this way is XLA should just work out of the box, and you wouldn't need to support a new op.

It seems like `tf.gather_nd` will only gather slices horizontally, so I had to transpose `mat_b` to extract the columns. From there instead of `tf.matmul` I just did an element-wise multiplication and then used `tf.math.reduce_sum` to get the dot product. But I think you're right in that we can do this without a new op. I'll see about just implementing it in `math_ops.py`.

```
>>> import tensorflow as tf
2024-01-05 13:56:56.117075: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
>>> indices = tf.constant([[[0,1],[1,0]],[[0,0],[1,1]]])
2024-01-05 13:57:04.040611: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.096683: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.096895: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.098041: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.098207: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.098354: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.170050: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.170255: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.170419: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:962] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-05 13:57:04.170549: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 5734 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3070, pci bus id: 0000:2d:00.0, compute capability: 8.6
>>> values = tf.constant([[0.12,0.0256],[0.54,0.77]])
>>> mat1 = tf.constant([[[0.7,0.094,1.043],[1.4,0.26,6.03]],[[0.842,0.439,0.218],[0.666,0.988,0.407]]])
>>> mat2 = tf.constant([[[0.076,0.99],[1.23,1.04],[0.98,0.326]],[[0.309,0.208],[0.557,0.4681],[0.6,0.742]]])
>>> alpha = 0.43
>>> beta = 0.77
>>> indices_dims = indices.get_shape().ndims
>>> row_indices = tf.slice(indices, [0] * indices_dims, [-1] * (indices_dims - 1) + [1])
>>> col_indices = tf.slice(indices, [0] * (indices_dims - 1) + [1], [-1] * (indices_dims - 1) + [1])
>>> rows = tf.gather_nd(mat1, row_indices, batch_dims=indices_dims-2)
>>> mat2_transpose = tf.transpose(mat2, perm=[0] * (indices_dims - 2) + [indices_dims - 1] + [indices_dims - 2])
>>> cols = tf.gather_nd(mat2_transpose, col_indices, batch_dims=indices_dims - 2)
>>> product = rows * cols
>>> dot = tf.reduce_sum(product, axis=-1)
>>> res = alpha * dot + beta * values
>>> print(res)
tf.Tensor(
[[0.57863456 2.7440202 ]
 [0.68906546 0.981192  ]], shape=(2, 2), dtype=float32)
```
Relates to #56311 

I have the operation implemented for cpu/gpu, but not for tpu/xla. If someone has time, I'd appreciate a quick review to make sure I have the right idea as well as any advice for where to go from here. Currently I only have the op registered for the float type.