title: Python: How to type hint tf.keras object in functions?
body: <p>This example function returns a dictionary of keras tensors:</p>
  <pre><code>import pandas as pd
  import tensorflow as tf
  
  def create_input_tensors(data: pd.DataFrame) -&gt; Dict[str,tf.keras.engine.keras_tensor.KerasTensor]:
     &quot;&quot;&quot;Turns each dataframe column into a keras tensor and returns them as a dict&quot;&quot;&quot;
     tensors = {}
     for name, column in data.items():
        tensors[name] = tf.keras.Input(shape=(1, ), name=name, dtype=float32)
     return tensors
  </code></pre>
  <p>I do not know how to correctly type hint the return value. Running the code snippet yields the following exception:</p>
  <pre><code>Exception has occurred: AttributeError
  module 'keras.api._v2.keras' has no attribute 'engine'
  </code></pre>
  <p>Googling this exception did not help. Running <code>type(tensors['year'])</code> in the debugger at the end of the function to see what type one of the elements in the return dictionary is (<code>year</code> is one of the columns in <code>data</code>) yields <code>&lt;class 'keras.engine.keras_tensor.KerasTensor'&gt;</code>.</p>
  <p>I have issues with this specific function, but also generally when trying to type hint functions that handle any kind of keras object. An answer that is applicable to these similar problems would be much appreciated.</p>