| 76 | |
| 77 | |
| 78 | class MockOperator: |
| 79 | def __init__(self, batch_size=None, addend=None): |
| 80 | self._batch_size = batch_size |
| 81 | self._addend = addend |
| 82 | |
| 83 | def _run(self, ctx, *inputs, batch_size=None, addend=None): |
| 84 | if addend is None: |
| 85 | addend = self._addend |
| 86 | if addend is None: |
| 87 | addend = 1 |
| 88 | if batch_size is None: |
| 89 | batch_size = self._batch_size |
| 90 | if len(inputs) == 0: |
| 91 | if isinstance(addend, MockBatch): |
| 92 | assert batch_size is not None |
| 93 | tensors = addend.tensors |
| 94 | else: |
| 95 | if not isinstance(addend, MockTensor): |
| 96 | addend = MockTensor(addend) |
| 97 | tensors = [addend] * (batch_size or 1) |
| 98 | return (tensors[0],) if batch_size is None else (MockBatch(tensors),) |
| 99 | if batch_size is not None: |
| 100 | return tuple( |
| 101 | MockBatch( |
| 102 | [ |
| 103 | MockTensor(t.data + a.data, t.layout) |
| 104 | for t, a in zip(input.tensors, addend.tensors) |
| 105 | ] |
| 106 | ) |
| 107 | for input in inputs |
| 108 | ) |
| 109 | else: |
| 110 | return tuple(MockTensor(t.data + addend.data, t.layout) for t in inputs) |
| 111 | |
| 112 | def _infer_num_outputs(self, *inputs, **args): |
| 113 | return max(len(inputs), 1) |
| 114 | |
| 115 | def _infer_output_devices(self, *inputs, **args): |
| 116 | return [input.device for input in inputs] |
| 117 | |
| 118 | |
| 119 | def test_mock_operator_tensor(): |
no outgoing calls