A check fail can be triggered in MatrixSquareRoot
Hi @shijy16,
Thank you for reporting the issue!
I was able to replicate the issue in Ubuntu 20.04. Please find the screenshot below. We are trying to dig deep into the issue, we'll update here soon. 
![image](https://user-images.githubusercontent.com/98147397/226864853-69cfb444-8b9c-471f-bea5-d797af6165f8.png)
Thank you!     I too replicated the same error as reported.
```
(base) suryanarayanay@surya-ubuntu-22-04:~$ python 60035.py
2023-03-24 11:56:04.828135: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-03-24 11:56:04.888753: 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: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-24 11:56:05.854504: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2.13.0-dev20230324
2.13.0-dev20230324
2023-03-24 11:56:13.188061: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 38197 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:00:04.0, compute capability: 8.0
2023-03-24 11:56:13.189779: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 38197 MB memory:  -> device: 1, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:00:05.0, compute capability: 8.0
2023-03-24 11:56:14.020498: F tensorflow/core/framework/tensor_shape.cc:201] Non-OK-status: InitDims(dim_sizes) status: INVALID_ARGUMENT: Encountered overflow when multiplying 4 with 5849693008847355793, result: -1
Aborted (core dumped)
(base) suryanarayanay@surya-ubuntu-22-04:~$ 
``` From the API document,   https://www.tensorflow.org/api_docs/python/tf/raw_ops/MatrixSquareRoot,  "The input is a tensor of shape [..., M, M] whose inner-most 2 dimensions form square matrices. ",

The input does not form a square matries for inner-most 2 dimensions.  

@shijy16 ,  could you share more background info why you choose such shape input ?
Thanks  @maxwillzq 
Hi, I am a security researcher, and I was testing the TensorFlow API to ensure its security. Although the input shape may not be the expected input according to the documentation or the code, it can pose a security risk, as it can lead to a crash. Therefore, this kind of input need to be checked and reported with proper error message, instead of throwing a crash.

I hope this answers your question. If you have any further inquiries or concerns, please do not hesitate to let me know. Thanks a lot for the explain, @shijy16.  It makes sense. I will take a look how to add this check. @shijy16 Please don't file vulnerabilities on GitHub. Please consult https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md and follow rules for responsible disclosure. Hi @shijy16 ,

The issue got resolved now with tf-nightly(2.14.0-dev20230514) and now TF is successfully able to raise the intended Error.
Please refer to attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/b1ca366360d589d5167da9587b9e301e/60035_nightly-2-14-0-dev20230514.ipynb).

Please check and confirm and let us know if we can close the issue as it resolved now.

Thanks! 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/60035">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/60035">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

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

Yes

### Source

binary

### Tensorflow Version

tf 2.13.0-dev20230317

### Custom Code

No

### OS Platform and Distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

CUDA 11.5

### GPU model and memory

_No response_

### Current Behaviour?

```shell
A crash due to check fail can be triggered by feeding a corner case to `tf.raw_ops.MatrixSquareRoot`. This case is similar to the one I submitted in the last issue, and this kind of bug is likely to be caused by the 0 and large value in input's shape.
```


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
input = tf.random.uniform([10, 15, 0, 4, 5849693008847355793], dtype=tf.float64, minval=-1024, maxval=1024)
res = tf.raw_ops.MatrixSquareRoot(
    input=input,
)
```


### Relevant log output

```shell
2023-03-18 18:27:16.712528: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-03-18 18:27:16.760165: 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: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-18 18:27:17.510832: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2023-03-18 18:27:19.128040: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14577 MB memory:  -> device: 0, name: Tesla V100-PCIE-16GB, pci bus id: 0000:2f:00.0, compute capability: 7.0
2023-03-18 18:27:19.389933: F tensorflow/core/framework/tensor_shape.cc:201] Non-OK-status: InitDims(dim_sizes) status: INVALID_ARGUMENT: Encountered overflow when multiplying 4 with 5849693008847355793, result: -1
Aborted (core dumped)
```
</details>