tensor lost shape info after tf.math.unsorted_segment_*

[sonogoto]

Please make sure that this is a bug. As per our
GitHub Policy,
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 16.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: n/a
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.6.0
• Python version: Python 3.8.5
• Bazel version (if compiling from source): n/a
• GCC/Compiler version (if compiling from source): n/a
• CUDA/cuDNN version: 11.1/n/a
• GPU model and memory: GeForce GTX 1070 8117MiB

You can collect some of this information using our environment capture
script
You can also obtain the TensorFlow version with:

1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior

Describe the expected behavior

Contributing

• Do you want to contribute a PR? (yes/no): no
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.

import tensorflow as tf
import numpy as np


class M1(tf.keras.Model):
    def __init__(self, **kwargs):
        super(M1, self).__init__(**kwargs)
        self.fc = tf.keras.layers.Dense(units=1)

    @tf.function
    def call(self, inputs, training=None):
        x, y, z = inputs
        x = tf.math.unsorted_segment_sum(x, tf.squeeze(y), z)
        return self.fc(x)


class M2(tf.keras.Model):
    def __init__(self, **kwargs):
        super(M2, self).__init__(**kwargs)
        self.fc = tf.keras.layers.Dense(units=1)

    @tf.function
    def call(self, inputs, training=None):
        x, y, z = inputs
        x = tf.math.segment_sum(x, tf.squeeze(y))
        return self.fc(x)


def gen():
    for _ in range(1024):
        offset = np.random.randint(1, 10, size=1024)
        y = np.repeat(np.arange(1024), offset)
        z = 1024
        x = np.random.rand(offset.sum(), 32)
        yield (x, y, z), np.random.rand(1024)


ds = tf.data.Dataset.from_generator(
    gen,
    output_signature=((tf.TensorSpec(shape=(None, 32), dtype=tf.float64),
                      tf.TensorSpec(shape=(None, ), dtype=tf.int64),
                      tf.TensorSpec(shape=[], dtype=tf.int32)),
                       tf.TensorSpec(shape=(None, ), dtype=tf.float64))
)



m1, m2 = M1(), M2()
m1.compile(loss=tf.keras.losses.MeanSquaredError(),
           optimizer=tf.keras.optimizers.Adagrad())
m2.compile(loss=tf.keras.losses.MeanSquaredError(),
           optimizer=tf.keras.optimizers.Adagrad())
# tf.math.segment_* is OK
m2.fit(ds, epochs=8)
# tf.math.unsorted_segment_* FAILED
m1.fit(ds, epochs=8)

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

Traceback (most recent call last):
File "/tmp/pycharm_project_124/python/test/dev.py", line 54, in
m1.fit(ds, epochs=8)
File "/usr/local/miniconda3/lib/python3.8/site-packages/keras/engine/training.py", line 1184, in fit
tmp_logs = self.train_function(iterator)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 885, in call
result = self._call(*args, **kwds)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 933, in _call
self._initialize(args, kwds, add_initializers_to=initializers)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 759, in _initialize
self._stateful_fn._get_concrete_function_internal_garbage_collected( # pylint: disable=protected-access
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3066, in _get_concrete_function_internal_garbage_collected
graph_function, _ = self._maybe_define_function(args, kwargs)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3463, in _maybe_define_function
graph_function = self._create_graph_function(args, kwargs)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3298, in _create_graph_function
func_graph_module.func_graph_from_py_func(
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1007, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 668, in wrapped_fn
out = weak_wrapped_fn().wrapped(*args, **kwds)
File "/usr/local/miniconda3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 994, in wrapper
raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

/usr/local/miniconda3/lib/python3.8/site-packages/keras/engine/training.py:853 train_function  *
    return step_function(self, iterator)
/tmp/pycharm_project_124/python/test/dev.py:14 call  *
    return self.fc(x)
/usr/local/miniconda3/lib/python3.8/site-packages/keras/engine/base_layer.py:1030 __call__  **
    self._maybe_build(inputs)
/usr/local/miniconda3/lib/python3.8/site-packages/keras/engine/base_layer.py:2659 _maybe_build
    self.build(input_shapes)  # pylint:disable=not-callable
/usr/local/miniconda3/lib/python3.8/site-packages/keras/layers/core.py:1175 build
    raise ValueError('The last dimension of the inputs to `Dense` '

ValueError: The last dimension of the inputs to `Dense` should be defined. Found `None`.

Process finished with exit code 1

[Saduf2019]

@sonogoto
Please refer to similar error issue and let us know:link,link2

[sonogoto]

@sonogoto Please refer to similar error issue and let us know:link,link2

I do not think it is the same issue with link2 since there is no broadcast operation in my code at all.

[Saduf2019]

@sonogoto
Does link1 help.

[sonogoto]

@sonogoto Does link1 help.

I'm afraid not

[sachinprasadhs]

Was able to reproduce your issue in Tensorflow 2.7, please find the gist here. Thanks!

[sonogoto]

Was able to reproduce your issue in Tensorflow 2.7, please find the gist here. Thanks!

It was