def test_forward_trace_functor():
from nncf.dynamic_graph.patch_pytorch import ForwardTraceOnly
from nncf.dynamic_graph.trace_tensor import TracedTensor, TensorMeta
func = ForwardTraceOnly()
shape1, shape2 = ([32, 1, 4, 8], [1, 8, 12, 16])
meta1, meta2 = (TensorMeta(5, 1, shape1), TensorMeta(3, 8, shape2))
input_tensor1 = TracedTensor.from_torch_tensor(torch.Tensor(shape1), meta1)
input_tensor2 = TracedTensor.from_torch_tensor(torch.Tensor(shape2), meta2)
# 1 -> 1
output_tensor = func(torch.Tensor.view, input_tensor1, [-1])
assert output_tensor.tensor_meta == input_tensor1.tensor_meta
# 1 -> N
outputs = func(torch.Tensor.chunk, input_tensor1, 3)
for out in outputs:
assert out.tensor_meta == input_tensor1.tensor_meta
# N -> N (2 -> 2)
outputs = func(lambda x: x + [5], [input_tensor1, input_tensor2])
assert outputs[0].tensor_meta == input_tensor1.tensor_meta
assert outputs[1].tensor_meta == input_tensor2.tensor_meta
# M -> N (2 -> 3)
with pytest.raises(RuntimeError):
outputs = func(lambda x: x + [torch.Tensor(shape2)],
[input_tensor1, input_tensor2])
# M -> N (2 -> 1)
with pytest.raises(RuntimeError):
outputs = func(lambda x: x[0], [input_tensor1, input_tensor2])
def make_op_exec_context_for_coalescing_test(scope_str: str) -> OperationExecutionContext:
ia_op_exec_context = InputAgnosticOperationExecutionContext.from_str(scope_str)
op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name,
ia_op_exec_context.scope_in_model,
ia_op_exec_context.call_order,
[TensorMeta(0, 0, [1])])
return op_exec_context
def __call__(self, lhs: TensorMeta, rhs: TensorMeta) -> bool:
return TensorMeta.default_comparator(lhs, rhs)