Rate this Page
★★★★★

python.assert#

dynamic_shape_assert#

Note

Tags:python.assert

Support Level: SUPPORTED

Original source code:

# mypy: allow-untyped-defsimporttorchclassDynamicShapeAssert(torch.nn.Module):"""    A basic usage of python assertion.    """defforward(self,x):# assertion with error messageassertx.shape[0]>2,f"{x.shape[0]} is greater than 2"# assertion without error messageassertx.shape[0]>1returnxexample_args=(torch.randn(3,2),)tags={"python.assert"}model=DynamicShapeAssert()torch.export.export(model,example_args)

Result:

ExportedProgram:classGraphModule(torch.nn.Module):defforward(self,x:"f32[3, 2]"):return(x,)Graphsignature:# inputsx:USER_INPUT# outputsx:USER_OUTPUTRangeconstraints:{}

list_contains#

Original source code:

# mypy: allow-untyped-defsimporttorchclassListContains(torch.nn.Module):"""    List containment relation can be checked on a dynamic shape or constants.    """defforward(self,x):assertx.size(-1)in[6,2]assertx.size(0)notin[4,5,6]assert"monkey"notin["cow","pig"]returnx+xexample_args=(torch.randn(3,2),)tags={"torch.dynamic-shape","python.data-structure","python.assert"}model=ListContains()torch.export.export(model,example_args)

Result:

ExportedProgram:classGraphModule(torch.nn.Module):defforward(self,x:"f32[3, 2]"):add:"f32[3, 2]"=torch.ops.aten.add.Tensor(x,x);x=Nonereturn(add,)Graphsignature:# inputsx:USER_INPUT# outputsadd:USER_OUTPUTRangeconstraints:{}