(device, layout=None)
| 27 | |
| 28 | |
| 29 | def build_src_pipe(device, layout=None): |
| 30 | if layout is None: |
| 31 | layout = "XY" |
| 32 | batches = [ |
| 33 | [ |
| 34 | np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32), |
| 35 | np.array([[10, 20], [30, 40], [50, 60]], dtype=np.float32), |
| 36 | ], |
| 37 | [ |
| 38 | np.array([[9, 10], [11, 12]], dtype=np.float32), |
| 39 | np.array([[100, 200, 300, 400, 500]], dtype=np.float32), |
| 40 | ], |
| 41 | ] |
| 42 | |
| 43 | src_pipe = Pipeline(len(batches), 1, 0) |
| 44 | out_batches = fn.external_source(source=batches, device=device, cycle=True, layout=layout) |
| 45 | src_pipe.set_outputs(out_batches) |
| 46 | return src_pipe, len(batches) |
| 47 | |
| 48 | |
| 49 | def _test_feed_input(device, is_serialized): |
no test coverage detected