python.assert ================= dynamic_shape_assert ^^^^^^^^^^^^^^^^^^^^ .. note:: Tags: :doc:`python.assert ` Support Level: SUPPORTED Original source code: .. code-block:: python import torch def dynamic_shape_assert(x): """ A basic usage of python assertion. """ # assertion with error message assert x.shape[0] > 2, f"{x.shape[0]} is greater than 2" # assertion without error message assert x.shape[0] > 1 return x Result: .. code-block:: ExportedProgram: class GraphModule(torch.nn.Module): def forward(self, l_x_: "f32[3, 2]"): return (l_x_,) Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=, arg=TensorArgument(name='l_x_'), target=None)], output_specs=[OutputSpec(kind=, arg=TensorArgument(name='l_x_'), target=None)]) Range constraints: {} Equality constraints: [] list_contains ^^^^^^^^^^^^^ .. note:: Tags: :doc:`torch.dynamic-shape `, :doc:`python.data-structure `, :doc:`python.assert ` Support Level: SUPPORTED Original source code: .. code-block:: python import torch def list_contains(x): """ List containment relation can be checked on a dynamic shape or constants. """ assert x.size(-1) in [6, 2] assert x.size(0) not in [4, 5, 6] assert "monkey" not in ["cow", "pig"] return x + x Result: .. code-block:: ExportedProgram: class GraphModule(torch.nn.Module): def forward(self, l_x_: "f32[3, 2]"): add: "f32[3, 2]" = torch.ops.aten.add.Tensor(l_x_, l_x_); l_x_ = None return (add,) Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=, arg=TensorArgument(name='l_x_'), target=None)], output_specs=[OutputSpec(kind=, arg=TensorArgument(name='add'), target=None)]) Range constraints: {} Equality constraints: []