See `Dataset.batch()` for details.
(self, input_dataset, batch_size, padded_shapes, padding_values,
drop_remainder)
| 3431 | """A `Dataset` that batches and pads contiguous elements from its input.""" |
| 3432 | |
| 3433 | def __init__(self, input_dataset, batch_size, padded_shapes, padding_values, |
| 3434 | drop_remainder): |
| 3435 | """See `Dataset.batch()` for details.""" |
| 3436 | self._input_dataset = input_dataset |
| 3437 | if sparse.any_sparse(get_legacy_output_classes(input_dataset)): |
| 3438 | # TODO(b/63669786): support batching of sparse tensors |
| 3439 | raise TypeError( |
| 3440 | "Batching of padded sparse tensors is not currently supported") |
| 3441 | self._input_dataset = input_dataset |
| 3442 | self._batch_size = ops.convert_to_tensor( |
| 3443 | batch_size, dtype=dtypes.int64, name="batch_size") |
| 3444 | padding_values = ( |
| 3445 | padding_values |
| 3446 | if padding_values is not None else _default_padding(input_dataset)) |
| 3447 | |
| 3448 | input_shapes = get_legacy_output_shapes(input_dataset) |
| 3449 | flat_padded_shapes = nest.flatten_up_to(input_shapes, padded_shapes) |
| 3450 | |
| 3451 | flat_padded_shapes_as_tensors = [] |
| 3452 | |
| 3453 | for input_component_shape, padded_shape in zip( |
| 3454 | nest.flatten(input_shapes), flat_padded_shapes): |
| 3455 | flat_padded_shapes_as_tensors.append( |
| 3456 | _padded_shape_to_tensor(padded_shape, input_component_shape)) |
| 3457 | |
| 3458 | self._padded_shapes = nest.pack_sequence_as(input_shapes, |
| 3459 | flat_padded_shapes_as_tensors) |
| 3460 | |
| 3461 | self._padding_values = nest.map_structure_up_to( |
| 3462 | input_shapes, _padding_value_to_tensor, padding_values, |
| 3463 | get_legacy_output_types(input_dataset)) |
| 3464 | self._drop_remainder = ops.convert_to_tensor( |
| 3465 | drop_remainder, dtype=dtypes.bool, name="drop_remainder") |
| 3466 | |
| 3467 | def _padded_shape_to_batch_shape(s): |
| 3468 | return tensor_shape.TensorShape([ |
| 3469 | tensor_util.constant_value(self._batch_size) |
| 3470 | if smart_cond.smart_constant_value(self._drop_remainder) else None |
| 3471 | ]).concatenate(tensor_util.constant_value_as_shape(s)) |
| 3472 | |
| 3473 | output_shapes = nest.map_structure( |
| 3474 | _padded_shape_to_batch_shape, self._padded_shapes) |
| 3475 | self._structure = structure.convert_legacy_structure( |
| 3476 | get_legacy_output_types(self._input_dataset), output_shapes, |
| 3477 | get_legacy_output_classes(self._input_dataset)) |
| 3478 | |
| 3479 | # pylint: disable=protected-access |
| 3480 | # TODO(jsimsa): Switch to using v2 only any time after 6/30/2018. |
| 3481 | if smart_cond.smart_constant_value(self._drop_remainder) is False: |
| 3482 | variant_tensor = gen_dataset_ops.padded_batch_dataset( |
| 3483 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3484 | batch_size=self._batch_size, |
| 3485 | padded_shapes=[ |
| 3486 | ops.convert_to_tensor(s, dtype=dtypes.int64) |
| 3487 | for s in nest.flatten(self._padded_shapes) |
| 3488 | ], |
| 3489 | padding_values=nest.flatten(self._padding_values), |
| 3490 | output_shapes=structure.get_flat_tensor_shapes(self._structure)) |
nothing calls this directly
no test coverage detected