Function to run one iteration with one input.
(ctx, inputs)
| 95 | def __call__(self): |
| 96 | with self._distribution.scope(): |
| 97 | def step_fn(ctx, inputs): |
| 98 | """Function to run one iteration with one input.""" |
| 99 | gradients_fn = backprop.implicit_grad(self._loss_fn) |
| 100 | gradients_fn = optimizer_lib.get_filtered_grad_fn(gradients_fn) |
| 101 | |
| 102 | grads_and_vars = self.distribution.extended.call_for_each_replica( |
| 103 | gradients_fn, args=(ctx, inputs)) |
| 104 | # If threads use layers, then we need to run the first step |
| 105 | # sequentially, so that layers.build() is not executed in parallel. |
| 106 | # Otherwise, multiple sets of mirrored variables are going to be |
| 107 | # created. |
| 108 | return self._optimizer._distributed_apply( # pylint: disable=protected-access |
| 109 | self.distribution, grads_and_vars) |
| 110 | |
| 111 | # TODO(priyag): Return the outputs, context, etc as well. |
| 112 | ctx = self.distribution.extended.experimental_run_steps_on_iterator( |
nothing calls this directly
no test coverage detected