• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Specifying return data type for `py_func` calls
2
3The `py_func` op requires specifying a
4[data type](https://www.tensorflow.org/guide/tensors#data_types).
5
6When wrapping a function with `py_func`, for instance using
7`@autograph.do_not_convert(run_as=autograph.RunMode.PY_FUNC)`, you have two
8options to specify the returned data type:
9
10 * explicitly, with a specified `tf.DType` value
11 * by matching the data type of an input argument, which is then assumed to be
12     a `Tensor`
13
14Examples:
15
16Specify an explicit data type:
17
18```
19  def foo(a):
20    return a + 1
21
22  autograph.util.wrap_py_func(f, return_dtypes=[tf.float32])
23```
24
25Match the data type of the first argument:
26
27```
28  def foo(a):
29    return a + 1
30
31  autograph.util.wrap_py_func(
32      f, return_dtypes=[autograph.utils.py_func.MatchDType(0)])
33```
34