get_static_value does not support tf.identity

[ngc92]

System information

• TensorFlow version (you are using): 2.1
• Are you willing to contribute it (Yes/No): Yes

Describe the feature and the current behavior/state.
get_static_value currently cannot see through a tf.identityop, see the following minimal example

>>> @tf.function
... def test():
...     a = tf.constant([5, 4])
...     print(tf.get_static_value(a))
...     b = tf.identity(a)
...     print(tf.get_static_value(b))
...     
>>> test()
[5 4]
None

Will this change the current api? How?
No

Who will benefit with this feature?
I noticed this problem while investigating the reason for tf.sparse.retain losing shape information, even though by definition it should return the same shape as its input. Turns out it constructs a new SparseTensor with dense_shape given by tf.identity(input.dense_shape). There might be other sparse functions with similar problems, but I have not investigated.

I think improving get_static_shape is preferable to providing better shape info for sparse.retain, so for now I am not going to submit a feature request for that.

Any Other info.

[ngc92]

this should be fixed with the merging of #38006