I don’t know if you could minimize a share few lines of standalone code to reproduce this (or a Colab).
In the meantime you could check if it could be related to:
  

      github.com/tensorflow/tensorflow
  

  
    
  
	  
  

  
    
      [TF 2.4] KerasTensor breaks typing compatibility
    

    
      
        opened 
        
          
        
        Nov 5, 2020
      
      


      
        
          
          lgeiger
        
      
    

    
        
          type:bug
        
        
          comp:keras
        
        
          TF 2.4
        
    
  


  
    **System information**
- Have I written custom code (as opposed to using a stoc…k 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](https://github.com/tensorflow/community/blob/master/rfcs/20200211-tf-types.md) (/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:

```python
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`](https://github.com/tensorflow/tensorflow/blob/6d9e0887f6bce8fbeeb430364e520d05350d96d5/tensorflow/python/types/core.py#L40-L54). 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:
   ```python
   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.