element-wise multiplication overflow with large dimension tensors
Hi @yufang67, Apologies for the delay.
Thank you for reporting the issue. I was able to replicate the issue using Tensorflow 2.9.2 in Ubuntu 20.04. Please find the screenshot below. 
![image](https://user-images.githubusercontent.com/98147397/232728213-dca32a59-5c2b-43a2-84d3-045e312c1269.png)
 It seems like we have to dig deep into the issue, we'll update here soon. Thank you!  @synandi  Thanks.  for the tests, if you increases a little the dimension, you will get CUDA illegal memory. @yufang67, I was able to replicate the CUDA illegal memory, kindly find the screenshot below for the same.
![image](https://user-images.githubusercontent.com/98147397/233010166-ca2e3344-0bef-4571-a41f-f5fe2aa348b1.png)

 Could you please try again with the CUDA 11.2 and cuDNN 8.1 as these are the [tested build configurations](https://www.tensorflow.org/install/source#gpu) and let us know if the issue persists. Thank you!   Hi @synandi,
I tested on cuda11.2 and cuDNN8.1.1 and got
tensorflow.python.framework.errors_impl.InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:GPU:0 to /job:localhost/replica:0/task:0/device:CPU:0 in order to run Identity: Could not synchronize CUDA stream: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered [Op:Identity] Hi @synandi, any update on this issue ? thanks ! I can reproduce the CUDA_ERROR_ILLEGAL_ADDRESS error on a simpler example (I didn't test the value mismatch error since it was taking too long to run on my machine):

```
import tensorflow as tf

x = tf.constant(1., shape=(1, 1))
y = tf.tile(x, (2**28, 9))  # Number of elements cannot fit in int32 value
y.numpy()
z = y * 2
z.numpy()
```

The CUDA_ERROR_ILLEGAL_ADDRESS occurs for ops with MLIR-generated kernels like multiply and relu, but not for non-MLIR-generated kernels like relu6. So this is likely an issue with MLIR-generated kernels on tensors whose number of elements is overflowing an int32 value.

@akuegel, can you take a look?

 @kushanam was working on jit-based kernels with int64 index types. I see that only atanh was enabled as experimental kernel.
I am not sure what the plan forward is.
@frgossen or @sherhut might know. We added experimental support for some kernels. You can build Tensorflow at head with the experimental flag: `--//tensorflow/tensorflow/core/kernels/mlir_generated:enable_experimental`. We are in the process of enabling these kernels by default.  They are now enabled by default. See https://github.com/tensorflow/tensorflow/blob/master/RELEASE.md This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60330">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60330">No</a>
 Thanks @frgossen for the info. I tried the latest nightly version of tensorflow tf_nightly-2.14.0.dev20230706, i still get the CUDA_ERROR_ILLEGAL_ADDRESS from the example of reedwm. The experimental kernel will be used automatically ? how can i use  these experimental kernel ?
Thanks,
 The experimental kernels will be enabled if you build TF with `--//tensorflow/core/kernels/mlir_generated:enable_experimental` flag. The nightly will not have that. 

All the currently supported i64-indexing kernels are enabled by default though, so they are not experimental anymore. Do you know which kernel fails for you?  Thanks @frgossen . it failed on multiplication. 
So if i compile the current master branch with the mention flag, the kernels should support i64-indexing right ? I would think you do not need the flag at this point since the kernels are enabled by default. I will add a test for mul and see if I can reproduce this. What dtype is this?  I use this example:

```
import tensorflow as tf

x = tf.constant(1., shape=(1, 1))
y = tf.tile(x, (2**28, 9))  # Number of elements cannot fit in int32 value
y.numpy()
z = y * 2
z.numpy()

```
i suppose the default dtype is float32.

 I did a little extra testing and it turns out the threshold for "large tensors" is wrong. Thanks for finding this! I will land a fix soon. In the meantime, it will work with tensors > 2**32 elements  https://github.com/tensorflow/tensorflow/commit/cb2b5456ee5f74d5bacd96672db9251d519e1f02 hey @frgossen, could it solve https://github.com/tensorflow/tensorflow/issues/57950?
 Hi @frgossen , i confirmed that it works for multiplication now. Great thanks. 
However, i still found some ops does not work:
tf.math.softmax
tf.raw_ops.BiasAddGrad
tf.math.sigmoid

Could you have a look ? Thanks 64-bit indexing is not enabled for all kernels atm, including the cases you found. It requires a little bit of work to enable it for the fully JIT-compiled kernels of which there are more and more.  Hi @frgossen, 
is there any plan to fix these frequently used ops ?
Thanks  Sorry for the late reply. This is not a priority at the moment. Feel free to add support for more kernels though and send them my way. 
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.9.3

### Custom Code

Yes

### OS Platform and Distribution

ubuntu 20.04

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.6/8

### GPU model and memory

single A100 80G

### Current Behaviour?

tested with tensorflow==2.9.3 and numpy==1.24.2 on single A100 80G GPU. If use small memory GPU, you may get OOM before reproducing the issue.

when using dimension (524288, 16, 9, 32), get illegal memory.
when using dimension (524288, 16, 8, 32), get Mismatched elements: 1024 / 2147483648 (4.77e-05%)
when using dimension (524288, 16, 7, 32), get correct values.

same behavior on eager mode and graph mode.
note: one related issue has been reported https://github.com/keras-team/tf-keras/issues/124

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import numpy as np

def test_mul_eager(input_shape):
    rng = np.random.RandomState(42)

    grad = rng.exponential(size=input_shape).astype(np.float32)
    grad_loss = rng.exponential(size=(input_shape[0],1,1,1)).astype(np.float32)
    
    with tf.device('/GPU:0'):
        tf_grad = tf.convert_to_tensor(grad)
        tf_grad_loss = tf.convert_to_tensor(grad_loss)
        out = tf_grad * tf_grad_loss
        #tf.print("==== shape ", tf_grad.shape, tf_grad_loss.shape, out.shape)


    with tf.device('/CPU:0'):
        out_cpu = tf.identity(out)
        tf_grad_cpu = tf.identity(tf_grad)
        tf_grad_loss_cpu = tf.identity(tf_grad_loss)
    #
    np.testing.assert_allclose(grad, tf_grad_cpu.numpy(), rtol=1e-5, atol=1e-4)
    np.testing.assert_allclose(grad_loss, tf_grad_loss_cpu.numpy(), rtol=1e-5, atol=1e-4)
    np.testing.assert_allclose(grad*grad_loss, out_cpu.numpy(), rtol=1e-5, atol=1e-4)

@tf.function
def compute_mul(b,t,u,v):
    with tf.device('/CPU:0'):
        x = tf.random.normal((1,t,u,v), dtype=tf.float32)
        y = tf.random.normal((1,1,1,1), dtype=tf.float32)
        tf_grad = tf.tile(x, (b,1,1,1))
        tf_grad_loss = tf.tile(y, (b,1,1,1))
    with tf.device('/GPU:0'):
        out = tf_grad * tf_grad_loss
        #out = tf.raw_ops.Mul(x=tf_grad, y=tf_grad_loss)
        #out = tf.multiply(tf_grad, tf_grad_loss)
    with tf.device('/CPU:0'):
        out_cpu = tf.identity(out)
        tf_grad_cpu = tf.identity(tf_grad)
        tf_grad_loss_cpu = tf.identity(tf_grad_loss)
    
    return out_cpu, tf_grad_cpu, tf_grad_loss_cpu

def test_mul_graph(input_shape):
    b,t,u,v = input_shape
    out_cpu, tf_grad_cpu, tf_grad_loss_cpu = compute_mul(b,t,u,v)
    
    np.testing.assert_allclose(tf_grad_cpu.numpy()*tf_grad_loss_cpu.numpy(), out_cpu.numpy(), rtol=1e-5, atol=1e-4)

if __name__ == '__main__':
    #input_shape = (524288, 16, 7, 32) # pass <2^31
    input_shape = (524288, 16, 8, 32) # value mismatch at 2^31 
    #input_shape = (524288, 16, 9, 32) # illegal memory access
    
    #test_mul_eager(input_shape)
    test_mul_graph(input_shape)
```


### Relevant log output

```shell
when using (524288, 16, 8, 32), got
Traceback (most recent call last):
  File "multiplication_mismatch.py", line 93, in <module>
    test_mul_graph(input_shape)
  File "multiplication_mismatch.py", line 85, in test_mul_graph
    np.testing.assert_allclose(tf_grad_cpu.numpy()*tf_grad_loss_cpu.numpy(), out_cpu.numpy(), rtol=1e-5, atol=1e-4)
  File "/fs/scratch/work/yu_fang/warp-transducer/venv_tf_29/lib/python3.8/site-packages/numpy/testing/_private/utils.py", line 1592, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/usr/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/fs/scratch/work/yu_fang/warp-transducer/venv_tf_29/lib/python3.8/site-packages/numpy/testing/_private/utils.py", line 862, in assert_array_compare
    raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=1e-05, atol=0.0001

Mismatched elements: 1024 / 2147483648 (4.77e-05%)
Max absolute difference: 2.2970564
Max relative difference: 0.
 x: array([[[[ 1.605678, -0.261173, -1.222985, ..., -1.186496, -0.111071,
           0.792078],
         [ 0.307934,  0.016565, -0.576156, ..., -0.17745 , -0.993849,...
 y: array([[[[ 0.      ,  0.      ,  0.      , ...,  0.      ,  0.      ,
           0.      ],
         [ 0.      ,  0.      ,  0.      , ...,  0.      ,  0.      ,...
```
</details>