tensorflow: [TF 2.4] KerasTensor breaks typing compatibility

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): macOS
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): v2.4.0-rc0 and tf-nightly
• Python version: 3.8

The Problem

Since version 2.4 functional Keras models use KerasTensor instead of tf.Tensor as layer output types. Unfortunately KerasTensor doesn’t subclass tf.Tensor which breaks isinstance(x, tf.Tensor) checks: https://github.com/tensorflow/tensorflow/blob/6d9e0887f6bce8fbeeb430364e520d05350d96d5/tensorflow/python/keras/engine/keras_tensor.py#L63

The release notes recommend to use tf.is_tensor instead. In my opinion this is not really Pythonic and breaks compatibility with the TF Types RFC (/cc @mdanatg) wich even mentiones that tf.is_tensor is expected to be deprecated eventually.

Concretely, switching from isinstance(x, tf.Tensor) to tf.is_tensor is also not possible in all cases. E.g. this breaks usage of static typecheckers like pytype or typeguard:

A common pattern which can also be found in TensorFlow Addons (/cc @seanpmorgan) is the following:

from typeguard import typechecked
import tensorflow as tf

@typechecked
def foo(x: tf.Tensor):
    print(x.dtype)

foo(tf.keras.Input(shape=(32, 32, 3)))  # Throws in TF 2.4 since `isintance` is used for typechecking

Possible solutions

1. Make KerasTensor a subclass of tf.Tensor. @mihaimaruseac @fchollet is there a reason why this isn’t currently the case?

2. Make KerasTensor a subclass of types.Tensor. I don’t see any disadvantage of doing this in general, but it wouldn’t really fix this issue since types.Tensor is not exposed as part of the public API yet so users would need to rely on private TensorFlow APIs

3. In usercode this could be fixed by directly relying on KerasTensor to replace the usage of tf.Tensor with:

   from typing import Union
   from tensorflow.python.keras.engine.keras_tensor import KerasTensor
   
   TensorType = Union[tf.Tensor, KerasTensor]
   

   I do not think this is a proper solution since it requires users to rely on internal APIs and implementation details that might change in the future.

I am currious to hear back from you on what the bestpractices for type checking of Tensors are, or whether I am just missing somthing obvious here.