(*args)
| 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)]) |
| 90 | |
| 91 | with ops.name_scope("batch") as name: |
| 92 | for a in args: |
| 93 | if not isinstance(a, ops.Tensor): |
| 94 | raise ValueError("All arguments to functions decorated with " |
| 95 | "`batch_function` are supposed to be Tensors; " |
| 96 | "found %s" % repr(a)) |
| 97 | return gen_batch_ops.batch_function( |
| 98 | num_batch_threads=num_batch_threads, |
| 99 | max_batch_size=max_batch_size, |
| 100 | batch_timeout_micros=batch_timeout_micros, |
| 101 | allowed_batch_sizes=allowed_batch_sizes, |
| 102 | max_enqueued_batches=max_enqueued_batches, |
| 103 | shared_name=name, |
| 104 | f=computation, |
| 105 | in_tensors=list(args), |
| 106 | captured_tensors=computation.captured_inputs, |
| 107 | Tout=[o.dtype for o in computation.outputs]) |
| 108 | |
| 109 | return decorated |
| 110 |
nothing calls this directly
no test coverage detected