torch.jit.interface#

torch.jit.interface(obj)[source]#

Decorate to annotate classes or modules of different types.

This decorator can be used to define an interface that can be used to annotateclasses or modules of different types. This can be used for to annotate a submoduleor attribute class that could have different types that implement the sameinterface, or which could be swapped at runtime; or to store a list of modules orclasses of varying types.

It is sometimes used to implement “Callables” - functions or modules that implementan interface but whose implementations differ and which can be swapped out.

Example:.. testcode:

importtorchfromtypingimportList@torch.jit.interfaceclassInterfaceType:defrun(self,x:torch.Tensor)->torch.Tensor:pass# implements InterfaceType@torch.jit.scriptclassImpl1:defrun(self,x:torch.Tensor)->torch.Tensor:returnx.relu()classImpl2(torch.nn.Module):def__init__(self)->None:super().__init__()self.val=torch.rand(())@torch.jit.exportdefrun(self,x:torch.Tensor)->torch.Tensor:returnx+self.valdefuser_fn(impls:List[InterfaceType],idx:int,val:torch.Tensor)->torch.Tensor:returnimpls[idx].run(val)user_fn_jit=torch.jit.script(user_fn)impls=[Impl1(),torch.jit.script(Impl2())]val=torch.rand(4,4)user_fn_jit(impls,0,val)user_fn_jit(impls,1,val)