tf.sparse.reduce_max should update the static shape

[yoeldr]

Click to expand!

Issue Type

Feature Request

Source

source

Tensorflow Version

tf 2.9.1

Custom Code

No

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 Behaviour?

tf.sparse.reduce_max(..., output_is_sparse=True) does not fill the result's static shape, although it can easily be done (like in the output_is_sparse=False case).

Standalone code to reproduce the issue

@tf.function
def f(n):
  t = tf.sparse.eye(n)
  tf.print(t.shape)
  result = tf.sparse.reduce_max(t, axis=1, keepdims=True, output_is_sparse=True)
  tf.print(result.shape)


f(3)

Relevant log output

TensorShape([3, 3])
TensorShape(None)

[sushreebarsa]

@gadagashwini I was able to replicate the issue, please find the gist here. Thank you!