(default_name,
batch_fn,
batch_size,
enqueue_many,
labeled_tensors,
allow_smaller_final_batch,
name=None)
| 404 | tc.Collection(core.LabeledTensorLike), bool, |
| 405 | tc.Optional(string_types)) |
| 406 | def _batch_helper(default_name, |
| 407 | batch_fn, |
| 408 | batch_size, |
| 409 | enqueue_many, |
| 410 | labeled_tensors, |
| 411 | allow_smaller_final_batch, |
| 412 | name=None): |
| 413 | with ops.name_scope(name, default_name, labeled_tensors) as scope: |
| 414 | labeled_tensors = [ |
| 415 | core.convert_to_labeled_tensor(lt) for lt in labeled_tensors |
| 416 | ] |
| 417 | |
| 418 | batch_ops = batch_fn([t.tensor for t in labeled_tensors], scope) |
| 419 | # TODO(shoyer): Remove this when they sanitize the TF API. |
| 420 | if not isinstance(batch_ops, list): |
| 421 | assert isinstance(batch_ops, ops.Tensor) |
| 422 | batch_ops = [batch_ops] |
| 423 | |
| 424 | if allow_smaller_final_batch: |
| 425 | batch_size = None |
| 426 | |
| 427 | @tc.returns(core.Axes) |
| 428 | @tc.accepts(core.Axes) |
| 429 | def output_axes(axes): |
| 430 | if enqueue_many: |
| 431 | if 'batch' not in axes or list(axes.keys()).index('batch') != 0: |
| 432 | raise ValueError( |
| 433 | 'When enqueue_many is True, input tensors must have an axis ' |
| 434 | 'called "batch" as their first dimension, ' |
| 435 | 'but axes were %s' % axes) |
| 436 | culled_axes = axes.remove('batch') |
| 437 | return core.Axes([('batch', batch_size)] + list(culled_axes.values())) |
| 438 | else: |
| 439 | return core.Axes([('batch', batch_size)] + list(axes.values())) |
| 440 | |
| 441 | output_labeled_tensors = [] |
| 442 | for i, tensor in enumerate(batch_ops): |
| 443 | axes = output_axes(labeled_tensors[i].axes) |
| 444 | output_labeled_tensors.append(core.LabeledTensor(tensor, axes)) |
| 445 | |
| 446 | return output_labeled_tensors |
| 447 | |
| 448 | |
| 449 | @tc.returns(tc.List(core.LabeledTensor)) |
no test coverage detected