Batches the computation done by the decorated function. So, for example, in the following code ```python @batch_function(1, 2, 3) def layer(a): return tf.matmul(a, a) b = layer(w) ``` if more than one session.run call is simultaneously trying to compute `b` the values of `w`
(num_batch_threads,
max_batch_size,
batch_timeout_micros,
allowed_batch_sizes=None,
max_enqueued_batches=10,
autograph=True)
| 30 | |
| 31 | @tf_export("nondifferentiable_batch_function") |
| 32 | def batch_function(num_batch_threads, |
| 33 | max_batch_size, |
| 34 | batch_timeout_micros, |
| 35 | allowed_batch_sizes=None, |
| 36 | max_enqueued_batches=10, |
| 37 | autograph=True): |
| 38 | """Batches the computation done by the decorated function. |
| 39 | |
| 40 | So, for example, in the following code |
| 41 | |
| 42 | ```python |
| 43 | @batch_function(1, 2, 3) |
| 44 | def layer(a): |
| 45 | return tf.matmul(a, a) |
| 46 | |
| 47 | b = layer(w) |
| 48 | ``` |
| 49 | |
| 50 | if more than one session.run call is simultaneously trying to compute `b` |
| 51 | the values of `w` will be gathered, non-deterministically concatenated |
| 52 | along the first axis, and only one thread will run the computation. See the |
| 53 | documentation of the `Batch` op for more details. |
| 54 | |
| 55 | Assumes that all arguments of the decorated function are Tensors which will |
| 56 | be batched along their first dimension. |
| 57 | |
| 58 | SparseTensor is not supported. The return value of the decorated function |
| 59 | must be a Tensor or a list/tuple of Tensors. |
| 60 | |
| 61 | Args: |
| 62 | num_batch_threads: Number of scheduling threads for processing batches |
| 63 | of work. Determines the number of batches processed in parallel. |
| 64 | max_batch_size: Batch sizes will never be bigger than this. |
| 65 | batch_timeout_micros: Maximum number of microseconds to wait before |
| 66 | outputting an incomplete batch. |
| 67 | allowed_batch_sizes: Optional list of allowed batch sizes. If left empty, |
| 68 | does nothing. Otherwise, supplies a list of batch sizes, causing the op |
| 69 | to pad batches up to one of those sizes. The entries must increase |
| 70 | monotonically, and the final entry must equal max_batch_size. |
| 71 | max_enqueued_batches: The maximum depth of the batch queue. Defaults to 10. |
| 72 | autograph: Whether to use autograph to compile python and eager style code |
| 73 | for efficient graph-mode execution. |
| 74 | |
| 75 | Returns: |
| 76 | The decorated function will return the unbatched computation output Tensors. |
| 77 | """ |
| 78 | |
| 79 | def decorator(fn): # pylint: disable=missing-docstring |
| 80 | |
| 81 | def decorated(*args): # pylint: disable=missing-docstring |
| 82 | |
| 83 | @function.defun(autograph=autograph) |
| 84 | def computation(*computation_args): |
| 85 | return fn(*computation_args) |
| 86 | |
| 87 | computation = computation.get_concrete_function( |
| 88 | *[tensor_spec.TensorSpec(dtype=x.dtype, shape=x.shape, name=str(i)) |
| 89 | for i, x in enumerate(args)]) |