python type hint - can tensorflow data type be used?

user_4281353 0
Viewed 2428

Is it possible to use the Tensorflow data types tf.dtypes.DType such as tf.int32 in Python type hint?

from typing import (
    Union,
)
import tensorflow as tf
import numpy as np


def f(
    a: Union[tf.int32, tf.float32]  # <----
): 
    return a * 2


def g(a: Union[np.int32, np.float32]):
    return a * 2


def test_a():
    f(tf.cast(1.0, dtype=tf.float32))  # <----
    g(np.float32(1.0))                 # Numpy type has no issue

It causes the error below and wonder if this is possible.

python3.8/typing.py:149: in _type_check
    raise TypeError(f"{msg} Got {arg!r:.100}.")
E   TypeError: Union[arg, ...]: each arg must be a type. Got tf.int32.
tensorflow type-hinting python
1 Answers

I assume you would like your function to accept:

  • tf.float32
  • np.float32
  • float
  • tf.int32
  • np.int32
  • int

and always return, say, tf.float32. Not completely sure if this covers your use case, but I would put a broad type for your input argument and cast to the desired type in your function.

experimental_follow_type_hints can be used along with type annotations to improve performance by reducing the number of expensive graph retracings. For example, an argument annotated with tf.Tensor is converted to Tensor even when the input is a non-Tensor value.

from typing import TYPE_CHECKING
import tensorflow as tf
import numpy as np


@tf.function(experimental_follow_type_hints=True)
def foo(x: tf.Tensor) -> tf.float32:
    if x.dtype == tf.int32:
        x = tf.dtypes.cast(x, tf.float32)
    return x * 2

a = tf.cast(1.0, dtype=tf.float32)
b = tf.cast(1.0, dtype=tf.int32)

c = np.float32(1.0)
d = np.int32(1.0)

e = 1.0
f = 1

for var in [a, b, c, d, e, f]:
    print(f"input: {var},\tinput type: {type(var)},\toutput: {foo(var)}\toutput type: {type(foo(var))}")

if TYPE_CHECKING:
    reveal_locals()

Output of python3 stack66968102.py:

input: 1.0,     input type: <class 'tensorflow.python.framework.ops.EagerTensor'>,      output: 2.0     output dtype: <dtype: 'float32'>
input: 1,       input type: <class 'tensorflow.python.framework.ops.EagerTensor'>,      output: 2.0     output dtype: <dtype: 'float32'>
input: 1.0,     input type: <class 'numpy.float32'>,    output: 2.0     output dtype: <dtype: 'float32'>
input: 1,       input type: <class 'numpy.int32'>,      output: 2.0     output dtype: <dtype: 'float32'>
input: 1.0,     input type: <class 'float'>,    output: 2.0     output dtype: <dtype: 'float32'>
input: 1,       input type: <class 'int'>,      output: 2.0     output dtype: <dtype: 'float32'>

Output of mypy stack66968102.py --ignore-missing-imports:

stack66968102.py:27: note: Revealed local types are:
stack66968102.py:27: note:     a: Any
stack66968102.py:27: note:     b: Any
stack66968102.py:27: note:     c: numpy.floating[numpy.typing._32Bit*]
stack66968102.py:27: note:     d: numpy.signedinteger[numpy.typing._32Bit*]
stack66968102.py:27: note:     e: builtins.float
stack66968102.py:27: note:     f: builtins.int
stack66968102.py:27: note:     tf: Any
stack66968102.py:27: note:     var: Any
user_500650
Related