Passes `args` to `self._func`, which is executed eagerly.
(self, device, token, args)
| 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): |
| 120 | if t.dtype.is_floating: |
| 121 | tape.watch(t) |
| 122 | ret = self._func(*args) |
| 123 | # copy the returned tensors to the PyFunc op's device if necessary. |
| 124 | device_name = device |
| 125 | if device_name is None: |
| 126 | # "None" here means "CPU", from the nullptr convention with C++ device |
| 127 | # pointers. |
| 128 | device_name = "/job:localhost/replica:0/task:0/device:CPU:0" |
| 129 | with ops.device(device): |
| 130 | if isinstance(ret, (tuple, list)): |
| 131 | outputs = [ |
| 132 | _maybe_copy_to_context_device(self._convert(x, dtype=dtype), |
| 133 | device_name) |
| 134 | for (x, dtype) in zip(ret, self._out_dtypes) |
| 135 | ] |
| 136 | elif ret is None: |
| 137 | outputs = None |
| 138 | else: |
| 139 | outputs = _maybe_copy_to_context_device( |
| 140 | self._convert(ret, dtype=self._out_dtypes[0]), device_name) |
| 141 | tape_cache[compat.as_bytes(token)] = (tape, args, outputs) |
| 142 | return outputs |
| 143 | |
| 144 | |
| 145 | class FuncRegistry(object): |
nothing calls this directly
no test coverage detected