Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf v2.20.0-rc0-4-g72fbba3d20f 2.20.0

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

Linux Ubuntu 22.04

Python version

3.9

Bazel version

None

GCC/compiler version

None

CUDA/cuDNN version

None

GPU model and memory

None

Current behavior?

The doc of tf.nn.weighted_moments() shows its description as below:

See the repro below, I find a phenomenon that when keepdims==True, the repro runs well, while when keepdims==False, the repro will raise error.

I tried to locate the error but find that I can solve this by set axes to just list type (that is [0]), instead of tf.Tensor, but why?

See the definition here, which shows that axes is needed to be a 1-d tensor of int32 values:

I find the error that is probably related here, which is not the place that the error shown(tensorflow/python/ops/gen_array_ops.py):

The place the error shown seems not to exist in original codes.

In all, I think that when I transfer keepdim from True into False, the repro should run well.

I think it is somehow a bug but I can't solve it.

Thanks for noting!

Standalone code to reproduce the issue

import tensorflow as tf
data = tf.random.normal(shape=(1, 1000, 1000))
axes = tf.constant([0], dtype=tf.int32) # choice: [0] or tf.constant([0], dtype=tf.int32)
frequency_weights = tf.constant(([1.0] * 1000))
keepdims = False # choice: True or False
moments = tf.nn.weighted_moments(data, axes, frequency_weights, keepdims=keepdims)

Relevant log output

~/anaconda3/envs/tensorflow-newest/lib/python3.9/site-packages/tensorflow/python/ops/gen_array_ops.py in ?(input, axis, name, ctx)
  10727 def squeeze_eager_fallback(input: Annotated[Any, TV_Squeeze_T], axis, name, ctx) -> Annotated[Any, TV_Squeeze_T]:
  10728   if axis is None:
  10729     axis = []
  10730   if not isinstance(axis, (list, tuple)):
> 10731     raise TypeError(
  10732         "Expected list for 'axis' argument to "
  10733         "'squeeze' Op, not %r." % axis)
  10734   axis = [_execute.make_int(_i, "axis") for _i in axis]

TypeError: Expected list for 'axis' argument to 'squeeze' Op, not <tf.Tensor: shape=(1,), dtype=int32, numpy=array([0], dtype=int32)>.