tf.keras.layers.GlobalAveragePooling2D() freezes tf.distribute.MirroredStrategy()
It may related to those two issues : #62234 #61314 Thanks for the reply @edwardyehuang!

I had a look at the two issues.

In this case the hangs with the specific model architecture above, but runs smooth with the minor modification. And on a single gpu everything runs.  > Thanks for the reply @edwardyehuang!
> 
> I had a look at the two issues.
> 
> In this case the hangs with the specific model architecture above, but runs smooth with the minor modification. And on a single gpu everything runs.

In the above issues, TensorFlow got stuck when using XLA + multi-gpus. Hi **@matthewhelmi**,
It looks like this is a duplicate of issue #[18862](https://github.com/keras-team/keras/issues/18862). Can you please close this issue, since it is already being tracked there.

Thank you! Thanks @edwardyehuang! is there something I can try related to the XLA+multigpu issue? Hi @Venkat6871,

I was told on the other issue to open it here as they believed it was TensorFlow related.

Would you suggest I close the other one? > Hi @Venkat6871,
> 
> I was told on the other issue to open it here as they believed it was TensorFlow related.
> 
> Would you suggest I close the other one?

Yes, you can keep this open and close the other one. 
### Issue type

Bug

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

No

### Source

source

### TensorFlow version

2.13.1

### Custom code

No

### OS platform and distribution

NAME="Amazon Linux" VERSION="2" ID="amzn" ID_LIKE="centos rhel fedora"

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

NVIDIA T4 (4 GPU) - VRAM 64gb

### Current behavior?

I was preparing my own training script for multi-gpu support to achieve higher batch sizes using tf.distribute.MirroredStrategy().

Things were working initially but began freezing (seemingly out of the blue) on NVIDIA T4 (4 GPU) - VRAM 64gb, but still runs fine on NVIDIA A10G (4 GPU) - VRAM 96gb.

I used the following example to debug on NVIDIA T4 (4 GPU)  - VRAM 64gb:

[Custom training with tf.distribute.Strategy](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/distribute/custom_training.ipynb#scrollTo=F2VeZUWUj5S4)

I found that if I substitute their model (See link) with

```
model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(10, kernel_size=3, activation="relu"),
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(10, activation="relu")
   ])
```
the training freezes from the start (just like my own code, which is using a similar model).

My symptoms are similar to [Multi-GPU training not starting or freezing. GPUs at 100%](https://github.com/tensorflow/tensorflow/issues/54325)

It stops freezing if I replace GlobalAveragePooling2D() with a Flatten().
```
model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(10, kernel_size=3, activation="relu"),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(10, activation="relu")
 ])
```
I have also noticed freezing when introducing stride > 1 in the conv layer.

Everything runs fine on a single-gpu. Freezing only occurs when using >1 GPU.

TF v2.13.1

### Standalone code to reproduce the issue

```shell
https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/distribute/custom_training.ipynb#scrollTo=F2VeZUWUj5S4
```


### Relevant log output

```shell
INFO:tensorflow:Collective all_reduce tensors: 4 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Collective all_reduce tensors: 4 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
```
