MemoryError: std::bad_alloc when calling tf.squeeze with a floating tensor axis
Hello, I'm glad you posted this issue. I can try to give you a hand.

First thing first, I would like to better clarify the `tf.squeeze()` function. According to the [documentation](https://www.tensorflow.org/api_docs/python/tf/squeeze), the `tf.squeeze()` function is used to remove dimensions of size 1 from the shape of a tensor. For instance, if you have a tensor `t` with shape `[batch_num, 1, elem_num]`, then `tf.squeeze(t)` will return a tensor with shape `[batch_num, elem_num]`. You can also specify which axis (or axes) to squeeze by passing the `axis` argument.

However, the `axis` argument MUST be integers, not floats. This is because the axis represent the indices of the dimensions to squeeze, and indices can only be integers.

In your specific case, the `input` tensor has shape `(3,)`, not `(1, 3)`. Therefore, it doesn't have any dimension of size 1, and you cannot squeeze it at all. The axis argument is irrelevant. I hope this solves your problem. I wish you a wonderful day!

Post Scriptum: to better check yourself, you can change your code this way:
```import tensorflow as tf
input = tf.constant([1,2,3], dtype='float32')
out = tf.squeeze(input)

print(input.shape)
print(out.shape)```  Hi @dante-tech 

Thanks for your explanation. 
My further inspection finds that the issue is not related to the data type of axis. Instead, this issue occurs as long as I set the axis to be a tensor. For instance, even though I set the axis data type to integer, this issue still occurs:

```
import tensorflow as tf
input = tf.constant([[1,2,3]], dtype='float32')
axis = tf.constant(1, dtype='int32')
out = tf.squeeze(input,axis)
```

error message:
```
/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/gen_array_ops.py in squeeze(input, axis, name)
  10613   if tld.is_eager:
  10614     try:
> 10615       _result = pywrap_tfe.TFE_Py_FastPathExecute(
  10616         _ctx, "Squeeze", name, input, "squeeze_dims", axis)
  10617       return _result

MemoryError: std::bad_alloc
```


 Hi @drewshark ,

The behaviour is due to the value passed to axis argument. The axis argument expects it as int or list of ints as per [documentation](https://www.tensorflow.org/api_docs/python/tf/squeeze#:~:text=input%20to%20squeeze.-,axis,to%20%5B%5D.%20If%20specified%2C%20only%20squeezes%20the%20dimensions%20listed.%20The,-dimension%20index%20starts).

Yeah, this error is related to Axis argument, but not related to input shape. If you pass invalid shapes or shapes that don't have dim '1' then it will raise intended exception.  I have attached a [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/06d4e539682780b4326a4f75f654a29e/62504.ipynb#scrollTo=NeGfeG5Jnhfn) for reference.

I will submit a probable fix for this to type check and raise exception. Hello @drewshark, 

Yes, you are correct. Is there anything else I can help you with? 
### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

tf 2.14.0, tf-nightly

### Custom code

Yes

### OS platform and distribution

_No response_

### 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

_No response_

### Current behavior?

When calling tf.squeeze with axis to be floating tensor, it will raise memory error. MemoryError looks dangerous to me, a proper input argument handling may be better.


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
input = tf.constant([1,2,3], dtype='float32')
axis = tf.constant(1.0, dtype='float32')
out = tf.squeeze(input,axis)
```


### Relevant log output

```shell
/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/gen_array_ops.py in squeeze(input, axis, name)
  10614   if tld.is_eager:
  10615     try:
> 10616       _result = pywrap_tfe.TFE_Py_FastPathExecute(
  10617         _ctx, "Squeeze", name, input, "squeeze_dims", axis)
  10618       return _result

MemoryError: std::bad_alloc
```
