(self, element_structure, batch_size,
expected_batched_structure)
| 644 | }), |
| 645 | ) |
| 646 | def testBatch(self, element_structure, batch_size, |
| 647 | expected_batched_structure): |
| 648 | batched_structure = nest.map_structure( |
| 649 | lambda component_spec: component_spec._batch(batch_size), |
| 650 | element_structure) |
| 651 | self.assertEqual(batched_structure, expected_batched_structure) |
| 652 | |
| 653 | @parameterized.named_parameters( |
| 654 | ("Tensor", tensor_spec.TensorSpec( |