Builds a training loop for TPUs. The set of loop-carried tensors corresponds to `inputs`. Both `condition` and `body` take the current value of the loop-carried tensors. 'body' additionally takes a tuple of infeed from infeed_queue if infeed_queue is not None. `condition` must return a s
(condition, body, inputs=None, infeed_queue=None, name=None)
| 28 | |
| 29 | |
| 30 | def while_loop(condition, body, inputs=None, infeed_queue=None, name=None): |
| 31 | """Builds a training loop for TPUs. |
| 32 | |
| 33 | The set of loop-carried tensors corresponds to `inputs`. Both |
| 34 | `condition` and `body` take the current value of the loop-carried |
| 35 | tensors. 'body' additionally takes a tuple of infeed from |
| 36 | infeed_queue if infeed_queue is not None. `condition` must return a |
| 37 | single boolean value that determines whether iteration |
| 38 | continues. `body` must return an updated list of values for the |
| 39 | loop-carried tensors. |
| 40 | |
| 41 | Args: |
| 42 | condition: a Python function that builds the loop condition. |
| 43 | body: a Python function that builds the loop body. |
| 44 | inputs: a list of initial values passed into the training loop, or |
| 45 | None (equivalent to an empty list). |
| 46 | infeed_queue: if not None, the infeed queue from which to append a tuple |
| 47 | of arguments as inputs to condition. |
| 48 | name: (Deprecated) Does nothing. |
| 49 | |
| 50 | Returns: |
| 51 | The final values of the loop-carried tensors. |
| 52 | |
| 53 | Raises: |
| 54 | TypeError: if body or condition has the wrong signature. |
| 55 | """ |
| 56 | del name |
| 57 | # Converts inputs to Tensors. |
| 58 | inputs = [] if inputs is None else [ops.convert_to_tensor(x) for |
| 59 | x in inputs] |
| 60 | input_types = [x.dtype for x in inputs] |
| 61 | input_arity = len(inputs) |
| 62 | |
| 63 | body_arg_error = xla.check_function_argument_count( |
| 64 | body, input_arity, infeed_queue) |
| 65 | if body_arg_error is not None: |
| 66 | if infeed_queue is None: |
| 67 | raise TypeError( |
| 68 | "Supplied loop body function cannot be called with the specified " |
| 69 | "inputs. You specified %d inputs: %s, but the loop body needs %s" % ( |
| 70 | input_arity, str([i.name for i in inputs]), body_arg_error)) |
| 71 | else: |
| 72 | raise TypeError( |
| 73 | "Supplied loop body function cannot be called with the specified " |
| 74 | "inputs. You specified %d inputs: %s and %d additional inputs from " |
| 75 | "infeed, but the computation needs %s" % (input_arity, str( |
| 76 | [i.name for i in inputs]), infeed_queue.number_of_tuple_elements, |
| 77 | body_arg_error)) |
| 78 | condition_arg_error = xla.check_function_argument_count( |
| 79 | condition, input_arity, None) |
| 80 | if condition_arg_error is not None: |
| 81 | if infeed_queue is None: |
| 82 | raise TypeError( |
| 83 | "Supplied loop condition function cannot be called with the " |
| 84 | "specified inputs. You specified %d inputs: %s, but the loop " |
| 85 | "condition needs %s" % (input_arity, str([i.name for i in inputs]), |
| 86 | condition_arg_error)) |
| 87 | else: |
no test coverage detected