A wrapper for a function owned by an EagerPyFunc.
| 60 | |
| 61 | |
| 62 | class EagerFunc(object): |
| 63 | """A wrapper for a function owned by an EagerPyFunc.""" |
| 64 | |
| 65 | def __init__(self, func, Tout, is_grad_func): |
| 66 | """Constructs an EagerFunc. |
| 67 | |
| 68 | Args: |
| 69 | func: The function to wrap. |
| 70 | Tout: A list of datatypes for the output; an empty list if the output is |
| 71 | None. |
| 72 | is_grad_func: Whether this EagerFunc is the gradient of another |
| 73 | EagerPyFunc. |
| 74 | """ |
| 75 | self._func = func |
| 76 | self._out_dtypes = Tout |
| 77 | self._is_grad_func = is_grad_func |
| 78 | |
| 79 | def _convert(self, value, dtype): |
| 80 | """Converts `value` to a tensor of type `dtype`, with error checking. |
| 81 | |
| 82 | Args: |
| 83 | value: The tensor to convert. |
| 84 | dtype: The desired dtype. |
| 85 | |
| 86 | Returns: |
| 87 | A tensor of type `dtype`, or a zeros tensor if value is None and |
| 88 | this function is in fact a grdient function. |
| 89 | |
| 90 | Raises: |
| 91 | RuntimeError: if `value` is a variable. |
| 92 | """ |
| 93 | |
| 94 | if isinstance(value, resource_variable_ops.ResourceVariable): |
| 95 | raise RuntimeError( |
| 96 | "Attempting to return a variable from an eagerly executed py_func. " |
| 97 | "Only numeric data structures like Tensors or NumPy arrays should " |
| 98 | "be returned; to return the value of a variable, make sure to obtain " |
| 99 | "the Tensor backing it by calling `.read_value()` on the variable in " |
| 100 | "question: %s" % value) |
| 101 | if value is None and self._is_grad_func: |
| 102 | # Gradient functions may legitimately return a list that contains |
| 103 | # both Tensors and Python Nones. Unfortuantely this breaks the |
| 104 | # OpKernel, so for now we replace None objects with zeros, which is |
| 105 | # mathematically correct but will prevent short-circuiting gradient |
| 106 | # computations. |
| 107 | # |
| 108 | # TODO(akshayka): Make it possible to return a list of both Tensors and |
| 109 | # Nones from an EagerPyFunc. |
| 110 | return constant_op.constant(0.0, dtype=dtype) |
| 111 | return ops.convert_to_tensor(value, dtype=dtype) |
| 112 | |
| 113 | def __call__(self, device, token, args): |
| 114 | """Passes `args` to `self._func`, which is executed eagerly.""" |
| 115 | |
| 116 | with context.eager_mode(), backprop.GradientTape() as tape: |
| 117 | # Only watch tensors with a floating dtype. |
| 118 | for tensor in args: |
| 119 | for t in nest.flatten(tensor): |