(self)
| 296 | self.assertEqual(test_case['shape'][1:], inferred_shape[1:]) |
| 297 | |
| 298 | def testFlatMap(self): |
| 299 | test_cases = [{ |
| 300 | 'tensor': 0, |
| 301 | 'shape': tensor_shape.TensorShape([]) |
| 302 | }, { |
| 303 | 'tensor': np.array([1, 2, 3]), |
| 304 | 'shape': tensor_shape.TensorShape([3]) |
| 305 | }, { |
| 306 | 'tensor': np.array([[1, 2, 3]]), |
| 307 | 'shape': tensor_shape.TensorShape([1, 3]) |
| 308 | }] |
| 309 | |
| 310 | for test_case in test_cases: |
| 311 | with ops.Graph().as_default() as g: |
| 312 | dataset = dataset_ops.Dataset.range(42) |
| 313 | |
| 314 | def make_dataset(tensor): |
| 315 | |
| 316 | def dataset_fn(n): |
| 317 | return dataset_ops.Dataset.from_tensors(tensor).repeat(n) |
| 318 | |
| 319 | return dataset_fn |
| 320 | |
| 321 | dataset = dataset.flat_map(make_dataset(test_case['tensor'])) |
| 322 | iterator = dataset_ops.make_one_shot_iterator(dataset) |
| 323 | get_next = iterator.get_next() |
| 324 | train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP) |
| 325 | train_op.append(get_next) |
| 326 | mg = meta_graph.create_meta_graph_def(graph=g) |
| 327 | grappler_item = item.Item(mg) |
| 328 | op_properties = grappler_item.GetOpProperties() |
| 329 | self.assertEqual(test_case['shape'], |
| 330 | op_properties['IteratorGetNext'][0].shape) |
| 331 | |
| 332 | def testInterleave(self): |
| 333 | test_cases = [{ |
nothing calls this directly
no test coverage detected