(*args)
| 87 | """ |
| 88 | def decorator(f): # pylint: disable=missing-docstring |
| 89 | def decorated(*args): |
| 90 | with ops.name_scope("batch") as name: |
| 91 | for a in args: |
| 92 | if not isinstance(a, ops.Tensor): |
| 93 | raise ValueError("All arguments to functions decorated with " |
| 94 | "`batch_function` are supposed to be Tensors; " |
| 95 | "found %s" % repr(a)) |
| 96 | batched_tensors, batch_index, id_t = gen_batch_ops.batch( |
| 97 | args, |
| 98 | num_batch_threads=num_batch_threads, |
| 99 | max_batch_size=max_batch_size, |
| 100 | batch_timeout_micros=batch_timeout_micros, |
| 101 | max_enqueued_batches=max_enqueued_batches, |
| 102 | allowed_batch_sizes=allowed_batch_sizes, |
| 103 | grad_timeout_micros=grad_timeout_micros, |
| 104 | shared_name=name) |
| 105 | outputs = f(*batched_tensors) |
| 106 | if isinstance(outputs, ops.Tensor): |
| 107 | outputs_list = [outputs] |
| 108 | else: |
| 109 | outputs_list = outputs |
| 110 | with ops.name_scope("unbatch") as unbatch_name: |
| 111 | unbatched = [ |
| 112 | gen_batch_ops.unbatch(t, batch_index, id_t, |
| 113 | timeout_micros=unbatch_timeout_micros, |
| 114 | shared_name=unbatch_name + "/" + t.name) |
| 115 | for t in outputs_list] |
| 116 | if isinstance(outputs, ops.Tensor): |
| 117 | return unbatched[0] |
| 118 | return unbatched |
| 119 | return decorated |
| 120 | return decorator |
nothing calls this directly
no test coverage detected