tensorflow: Grappler error when Softmax input has a variable dimension

System information

• Have I written custom code: Yes
• OS Platform and Distribution: Ubuntu 18.04
• TensorFlow installed from: binary
• TensorFlow version: 2.5.0
• Python version: 3.6.9
• CUDA/cuDNN version: 11.2/8.1.0
• GPU model and memory: NVIDIA Quadro P620 (4GB)

Describe the current behavior

The grappler pass is logging an error when Softmax is used in a tf.function with one or more variable dimensions (e.g. a variable batch size). This log started appearing in TensorFlow 2.5.

It does not seem to cause issues when running the model, but it might indicate a bug in the grappler implementation.

Describe the expected behavior

This operation should not produce any warning or error.

Contributing

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

Standalone code to reproduce the issue

import tensorflow as tf

input_signature = (tf.TensorSpec([None, 20], tf.float32),)
softmax = tf.function(tf.nn.softmax, input_signature=input_signature)
softmax(tf.random.uniform([2, 20]))

Other info / logs

The code above logs the following error:

2021-07-02 07:34:12.478300: W tensorflow/core/grappler/costs/op_level_cost_estimator.cc:689] Error in PredictCost() for the op: op: "Softmax" attr { key: "T" value { type: DT_FLOAT } } inputs { dtype: DT_FLOAT shape { unknown_rank: true } } device { type: "GPU" vendor: "NVIDIA" model: "Quadro P620" frequency: 1442 num_cores: 4 environment { key: "architecture" value: "6.1" } environment { key: "cuda" value: "11020" } environment { key: "cudnn" value: "8100" } num_registers: 65536 l1_cache_size: 24576 l2_cache_size: 524288 shared_memory_size_per_multiprocessor: 98304 memory_size: 3092316160 bandwidth: 96128000 } outputs { dtype: DT_FLOAT shape { unknown_rank: true } }