ArrayContract

fromarraycontractimportshape, _importtorch@shape(x=(_, 'N'), y=('N', _))defmatrix_dot(x, y):
returnx @ ymatrix_dot(torch.rand(3,4), torch.rand(4,5)) # OKmatrix_dot(torch.rand(3,4), torch.rand(3,5)) # raise AssertionError

fromarraycontractimportshape, _importtorchfromtorchimportnnlinear=nn.Linear(3, 4)
@shape((..., 3))defforward_linear(x):
""" requires x.shape[-1] == 3 """returnlinear(x)
forward_linear(torch.rand(4,5,3)) # OKforward_linear(torch.rand(4,4)) # raise AssertionError

fromarraycontractimportdtypefromarraycontractimportndimimporttorch@ndim(x=3, y=4)defndim_contract(x, y):
print("requires x.ndim == 3 and y.ndim == 4")
@dtype(x=torch.long)defdtype_contract(x):
print("requires x.dtype == torch.long")

fromarraycontractimportTriggerfromarraycontractimportdtypeimporttorchTrigger.dtype_check_trigger=False@dtype(x=torch.long)defdtype_contract(x):
print("not requires x.dtype == torch.long")
dtype_contract(torch.rand(3, 4).float()) # OK