When to use#

In general, implement a custom function if you want to perform computations in your modelthat are not differentiable or rely on non-PyTorch libraries (e.g., NumPy), butstill wish for your operation to chain with other ops and work with the autograd engine.

In some situations, custom functions can also be used to improve performance andmemory usage: If you implemented your forward and backward passes using aC++ extension,you can wrap them inFunction to interface with the autogradengine. If you’d like to reduce the number of buffers saved for the backward pass,custom functions can be used to combine ops together.

When not to use#

If you can already write your function in terms of PyTorch’s built-in ops, itsbackward graph is (most likely) already able to be recorded by autograd. In this case, you donot need to implement the backward function yourself. Consider using a plainold Python function.

If you need to maintain state, i.e., trainable parameters, you should (also) use acustom module. See the section below for more information on extendingtorch.nn.

If you’d like to alter the gradients during the backward pass or perform a sideeffect, consider registering atensor orModule hook.

How to use#

Take the following steps:1. SubclassFunction and implement theforward(),(optional)setup_context() andbackward() methods.2. Call the proper methods on thectx argument.3. Declare whether your function supportsdouble backward.4. Validate whether your gradients are correct using gradcheck.

Step 1: After subclassingFunction, you’ll need to define 3 methods:

• forward() is the code that performs the operation. It can takeas many arguments as you want, with some of them being optional, if youspecify the default values. All kinds of Python objects are accepted here.Tensor arguments that track history (i.e., withrequires_grad=True) will be converted to ones that don’t track historybefore the call, and their use will be registered in the graph. Note that thislogic won’t traverse lists/dicts/any other data structures and will onlyconsider tensors that are direct arguments to the call. You canreturn either a singleTensor output, or atuple oftensors if there are multiple outputs. Also, please refer to thedocs ofFunction to find descriptions of useful methods that can becalled only fromforward().

• setup_context() (optional). One can either write a “combined”forward() thataccepts actx object or (as of PyTorch 2.0) a separateforward() that doesnot acceptctx and asetup_context() method where thectx modification happens.Theforward() should have the compute andsetup_context() shouldonly be responsible for thectx modification (and not have any compute).In general the separateforward() andsetup_context() is closer to howPyTorch native operations work and therefore more composable with various PyTorch subsystems.SeeCombined or separate forward() and setup_context() for more details.

• backward() (orvjp()) defines the gradient formula.It will be given as manyTensor arguments as there were outputs, with eachof them representing gradient w.r.t. that output. It is important NEVER to modifythese in-place. It should return as many tensors as therewere inputs, with each of them containing the gradient w.r.t. itscorresponding input. If your inputs didn’t require gradient(needs_input_grad is a tuple of booleans indicatingwhether each input needs gradient computation), or were non-Tensorobjects, you can returnpython:None. Also, if you have optionalarguments toforward() you can return more gradients than therewere inputs, as long as they’re allNone.

Step 2: It is your responsibility to use the functions inctxproperly in order to ensure that the newFunction works properly withthe autograd engine.

• save_for_backward() should beused to save any tensors needed for the backward pass (as opposed todirectly onctx). You cannot usesave_for_backward for non-tensors;you should store those directly onctx.

  Saving tensors viasave_for_backward:1. Allows the autograd engine to clearthem as soon as the backward computation of theautograd.Function completes.(If a tensor is stored directly onctxit will unnecessarily remain alive for the lifetime of the autograd graph –typically until the end of the iteration.)2. Helps avoid certain reference cycles, (e.g., since the tensoroutput of theautograd.Function itself keeps a reference to the ctx).3. Is important for compatibility withfeatures like activation checkpointing and offloading that rely ontorch.autograd.graph.saved_tensors_hooks.

  If tensors that are neither inputs nor outputs are saved for backward yourFunction may not support double backward (see step 3).

• mark_dirty() must be used tomark any input that is modified inplace by the forward function.

• mark_non_differentiable() mustbe used to tell the engine if an output is not differentiable. Bydefault all output tensors that are of differentiable type will be setto require gradient. Tensors of non-differentiable type (i.e., integral types)are never marked as requiring gradients.

• set_materialize_grads() can beused to tell the autograd engine to optimize gradient computations in the cases wherethe output does not depend on the input by not materializing grad tensors given to backwardfunction. That is, if set to False, None object in Python or “undefined tensor” (tensor x forwhich x.defined() is False) in C++ will not be converted to a tensor filled with zeros priorto calling backward, and so your code will need to handle such objects as if they weretensors filled with zeros. The default value of this setting is True.

Step 3: If yourFunction does not support double backwardyou should explicitly declare this by decorating backward with theonce_differentiable(). With this decorator, attempts toperform double backward through your function will produce an error.See our double backward tutorial for more information on double backward.

Step 4: It is recommended that you usetorch.autograd.gradcheck()to check whether your backward function correctly computes gradients of theforward by computing the Jacobian matrix using your backward function andcomparing the value element-wise with the Jacobian computed numerically usingfinite-differencing.

Example#

Below you can find code for aLinear function, withadditional comments:

# Inherit from FunctionclassLinearFunction(Function):# Note that forward, setup_context, and backward are @staticmethods@staticmethoddefforward(input,weight,bias):output=input.mm(weight.t())ifbiasisnotNone:output+=bias.unsqueeze(0).expand_as(output)returnoutput@staticmethod# inputs is a Tuple of all of the inputs passed to forward.# output is the output of the forward().defsetup_context(ctx,inputs,output):input,weight,bias=inputsctx.save_for_backward(input,weight,bias)# This function has only a single output, so it gets only one gradient@staticmethoddefbackward(ctx,grad_output):# This is a pattern that is very convenient - at the top of backward# unpack saved_tensors and initialize all gradients w.r.t. inputs to# None. Thanks to the fact that additional trailing Nones are# ignored, the return statement is simple even when the function has# optional inputs.input,weight,bias=ctx.saved_tensorsgrad_input=grad_weight=grad_bias=None# These needs_input_grad checks are optional and there only to# improve efficiency. If you want to make your code simpler, you can# skip them. Returning gradients for inputs that don't require it is# not an error.ifctx.needs_input_grad[0]:grad_input=grad_output.mm(weight)ifctx.needs_input_grad[1]:grad_weight=grad_output.t().mm(input)ifbiasisnotNoneandctx.needs_input_grad[2]:grad_bias=grad_output.sum(0)returngrad_input,grad_weight,grad_bias

Now, to make it easier to use these custom ops, we recommend either aliasingthem or wrapping them in a function. Wrapping in a function lets us supportdefault arguments and keyword arguments:

# Option 1: aliaslinear=LinearFunction.apply# Option 2: wrap in a function, to support default args and keyword args.deflinear(input,weight,bias=None):returnLinearFunction.apply(input,weight,bias)

Here, we give an additional example of a function that is parametrized bynon-Tensor arguments:

classMulConstant(Function):@staticmethoddefforward(tensor,constant):returntensor*constant@staticmethoddefsetup_context(ctx,inputs,output):# ctx is a context object that can be used to stash information# for backward computationtensor,constant=inputsctx.constant=constant@staticmethoddefbackward(ctx,grad_output):# We return as many input gradients as there were arguments.# Gradients of non-Tensor arguments to forward must be None.returngrad_output*ctx.constant,None

And here, we optimize the above example by calling set_materialize_grads(False):

classMulConstant(Function):@staticmethoddefforward(tensor,constant):returntensor*constant@staticmethoddefsetup_context(ctx,inputs,output):tensor,constant=inputsctx.set_materialize_grads(False)ctx.constant=constant@staticmethoddefbackward(ctx,grad_output):# Here we must handle None grad_output tensor. In this case we# can skip unnecessary computations and just return None.ifgrad_outputisNone:returnNone,None# We return as many input gradients as there were arguments.# Gradients of non-Tensor arguments to forward must be None.returngrad_output*ctx.constant,None

If you need any “intermediate” Tensors computed inforward() to be saved,either they must be returned as outputs, or combineforward andsetup_context()(seeCombined or separate forward() and setup_context()).Note that this means if you want gradients to flow through those intermediate values, youneed to define the gradient formula for them (see alsothe double backward tutorial):

classMyCube(torch.autograd.Function):@staticmethoddefforward(x):# We wish to save dx for backward. In order to do so, it must# be returned as an output.dx=3*x**2result=x**3returnresult,dx@staticmethoddefsetup_context(ctx,inputs,output):x,=inputsresult,dx=outputctx.save_for_backward(x,dx)@staticmethoddefbackward(ctx,grad_output,grad_dx):x,dx=ctx.saved_tensors# In order for the autograd.Function to work with higher-order# gradients, we must add the gradient contribution of `dx`,# which is grad_dx * 6 * x.result=grad_output*dx+grad_dx*6*xreturnresult# Wrap MyCube in a function so that it is clearer what the output isdefmy_cube(x):result,dx=MyCube.apply(x)returnresult

You probably want to check if the backward method you implemented actuallycomputes the derivatives of your function. It is possible by comparing withnumerical approximations using small finite differences:

fromtorch.autogradimportgradcheck# gradcheck takes a tuple of tensors as input, check if your gradient# evaluated with these tensors are close enough to numerical# approximations and returns True if they all verify this condition.input=(torch.randn(20,20,dtype=torch.double,requires_grad=True),torch.randn(30,20,dtype=torch.double,requires_grad=True))test=gradcheck(linear,input,eps=1e-6,atol=1e-4)print(test)

SeeNumerical gradient checking for more details on finite-difference gradient comparisons.If your function is used in higher order derivatives (differentiating the backward pass) youcan use thegradgradcheck function from the same package to check higher order derivatives.

Combined or separateforward() andsetup_context()#

There are two main ways to defineFunction. Either:

• define aforward() that combines the forward compute logic withsetup_context()

• (as of PyTorch 2.0) define a separateforward() andsetup_context()

We recommend the second option (separateforward() andsetup_context())because that is closer to how PyTorch native operations are implemented and it composeswithtorch.func transforms. However, we plan to support both approaches going forward;combiningforward() withsetup_context(): leads to more flexibility sinceyou are able to save intermediates without returning them as output.

Please see the previous section for how to defineFunction with separateforward() andsetup_context().

Here is an example of how to define aFunction with combinedforward() andsetup_context():

classLinearFunction(Function):@staticmethod# ctx is the first argument to forwarddefforward(ctx,input,weight,bias=None):# The forward pass can use ctx.ctx.save_for_backward(input,weight,bias)output=input.mm(weight.t())ifbiasisnotNone:output+=bias.unsqueeze(0).expand_as(output)returnoutput@staticmethoddefbackward(ctx,grad_output):input,weight,bias=ctx.saved_tensorsgrad_input=grad_weight=grad_bias=Noneifctx.needs_input_grad[0]:grad_input=grad_output.mm(weight)ifctx.needs_input_grad[1]:grad_weight=grad_output.t().mm(input)ifbiasisnotNoneandctx.needs_input_grad[2]:grad_bias=grad_output.sum(0)returngrad_input,grad_weight,grad_bias

Forward mode AD#

Overriding the forward mode AD formula has a very similar API with some different subtleties.You can implement thejvp() function.

It will be given as manyTensor arguments as there were inputs, with eachof them representing gradient w.r.t. that input. It should return as many tensors as therewere outputs, with each of them containing the gradient w.r.t. its corresponding output.Thejvp() will be called just after theforward()method, before theapply() returns.

jvp() has a few subtle differences with thebackward() function:

• You can use thectx to pass any data from theforward() to thejvp() function.If that state will not be needed for thebackward(),you can explicitly free it by doingdelctx.foo at the end of thejvp() function.

• The implementation ofjvp() must be backward differentiable or explicitly check thatnone of the given forward mode gradient hasrequires_grad set.

• Thejvp() function must match the view/inplace behavior offorward().For example, if thei th input is modified inplace, then thei th gradient must be updated inplace.Similarly, if thej th output is a view of thek th input. Then the returnedj th output gradient must bea view of the givenk th input gradient.

• Because the user cannot specify which gradient needs to be computed, thejvp() function shouldalways compute gradients for all the outputs.

• The forward mode gradients do respect the flag set byset_materialize_grads()and you can getNone input gradients when this is disabled.

torch.func transforms and/ortorch.vmap()#

Please seeExtending torch.func with autograd.Function for details.

Adding aModule#

Sincenn heavily utilizesautograd, adding a newModule requires implementing aFunctionthat performs the operation and can compute the gradient. From now on let’sassume that we want to implement aLinear module and we have the functionimplemented as in the listing above. There’s very little code required toadd this. Now, there are two functions that need to be implemented:

• __init__ (optional) - takes in arguments such as kernel sizes, numbersof features, etc. and initializes parameters and buffers.

• forward() - instantiates aFunction anduses it to perform the operation. It’s very similar to a functional wrappershown above.

This is how aLinear module can be implemented:

classLinear(nn.Module):def__init__(self,input_features,output_features,bias=True):super().__init__()self.input_features=input_featuresself.output_features=output_features# nn.Parameter is a special kind of Tensor, that will get# automatically registered as Module's parameter once it's assigned# as an attribute. Parameters and buffers need to be registered, or# they won't appear in .parameters() (doesn't apply to buffers), and# won't be converted when e.g. .cuda() is called. You can use# .register_buffer() to register buffers.# nn.Parameters require gradients by default.self.weight=nn.Parameter(torch.empty(output_features,input_features))ifbias:self.bias=nn.Parameter(torch.empty(output_features))else:# You should always register all possible parameters, but the# optional ones can be None if you want.self.register_parameter('bias',None)# Not a very smart way to initialize weightsnn.init.uniform_(self.weight,-0.1,0.1)ifself.biasisnotNone:nn.init.uniform_(self.bias,-0.1,0.1)defforward(self,input):# See the autograd section for explanation of what happens here.returnLinearFunction.apply(input,self.weight,self.bias)defextra_repr(self):# (Optional)Set the extra information about this module. You can test# it by printing an object of this class.return'input_features={}, output_features={}, bias={}'.format(self.input_features,self.output_features,self.biasisnotNone)

Extendingtorch with aTensor-like type#

To make this concrete, let’s begin with a simple example that illustrates theAPI dispatch mechanism. We’ll create a custom type that represents a 2D scalartensor, parametrized by the orderN and value along the diagonal entries,value:

classScalarTensor(object):def__init__(self,N,value):self._N=Nself._value=valuedef__repr__(self):return"ScalarTensor(N={}, value={})".format(self._N,self._value)deftensor(self):returnself._value*torch.eye(self._N)

This first iteration of the design isn’t very useful. The main functionality ofScalarTensor is to provide a more compact string representation of a scalartensor than in the base tensor class:

>>>d=ScalarTensor(5,2)>>>dScalarTensor(N=5, value=2)>>>d.tensor()tensor([[2., 0., 0., 0., 0.],        [0., 2., 0., 0., 0.],        [0., 0., 2., 0., 0.],        [0., 0., 0., 2., 0.],        [0., 0., 0., 0., 2.]])

If we try to use this object with thetorch API, we will runinto issues:

>>>importtorch>>>torch.mean(d)TypeError: mean(): argument 'input' (position 1) must be Tensor, not ScalarTensor

Adding a__torch_function__ implementation toScalarTensor makes itpossible for the above operation to succeed. Let’s re-do our implementation,this time adding a__torch_function__ implementation:

HANDLED_FUNCTIONS={}classScalarTensor(object):def__init__(self,N,value):self._N=Nself._value=valuedef__repr__(self):return"ScalarTensor(N={}, value={})".format(self._N,self._value)deftensor(self):returnself._value*torch.eye(self._N)@classmethoddef__torch_function__(cls,func,types,args=(),kwargs=None):ifkwargsisNone:kwargs={}iffuncnotinHANDLED_FUNCTIONSornotall(issubclass(t,(torch.Tensor,ScalarTensor))fortintypes):returnNotImplementedreturnHANDLED_FUNCTIONS[func](*args,**kwargs)

The__torch_function__ method takes four arguments:func, a referenceto the torch API function that is being overridden,types, the list oftypes of Tensor-likes that implement__torch_function__,args, thetuple of arguments passed to the function, andkwargs, the dict of keywordarguments passed to the function. It uses a global dispatch table namedHANDLED_FUNCTIONS to store custom implementations. The keys of thisdictionary are functions in thetorch namespace and the values areimplementations forScalarTensor.

This class definition isn’t quite enough to maketorch.mean do the rightthing when we pass it aScalarTensor – we also need to define animplementation fortorch.mean forScalarTensor operands and add theimplementation to theHANDLED_FUNCTIONS dispatch table dictionary. One wayof doing this is to define a decorator:

importfunctoolsdefimplements(torch_function):"""Register a torch function override for ScalarTensor"""defdecorator(func):functools.update_wrapper(func,torch_function)HANDLED_FUNCTIONS[torch_function]=funcreturnfuncreturndecorator

which can be applied to the implementation of our override:

@implements(torch.mean)defmean(input):returnfloat(input._value)/input._N

With this change we can now usetorch.mean withScalarTensor:

>>>d=ScalarTensor(5,2)>>>torch.mean(d)0.4

Of coursetorch.mean is an example of the simplest kind of function tooverride since it only takes one operand. We can use the same machinery tooverride a function that takes more than one operand, any one of which might bea tensor or tensor-like that defines__torch_function__, for example fortorch.add():

defensure_tensor(data):ifisinstance(data,ScalarTensor):returndata.tensor()returntorch.as_tensor(data)@implements(torch.add)defadd(input,other):try:ifinput._N==other._N:returnScalarTensor(input._N,input._value+other._value)else:raiseValueError("Shape mismatch!")exceptAttributeError:returntorch.add(ensure_tensor(input),ensure_tensor(other))

This version has a fast path for when both operands areScalarTensorinstances and also a slower path which degrades to converting the data totensors when either operand is not aScalarTensor. That makes the overridefunction correctly when either operand is aScalarTensor or a regularTensor:

>>>s=ScalarTensor(2,2)>>>torch.add(s,s)ScalarTensor(N=2, value=4)>>>t=torch.tensor([[1,1,],[1,1]])>>>torch.add(s,t)tensor([[3., 1.],        [1., 3.]])

Note that our implementation ofadd does not takealpha orout askeyword arguments liketorch.add() does:

>>>torch.add(s,s,alpha=2)TypeError: add() got an unexpected keyword argument 'alpha'

For speed and flexibility the__torch_function__ dispatch mechanism does notcheck that the signature of an override function matches the signature of thefunction being overridden in thetorch API. For some applications ignoringoptional arguments would be fine but to ensure full compatibility withTensor, user implementations of torch API functions should take care toexactly emulate the API of the function that is being overridden.

Functions in thetorch API that do not have explicit overrides willreturnNotImplemented from__torch_function__. If all operands with__torch_function__ defined on them returnNotImplemented, PyTorch willraise aTypeError. This means that most of the time operations that do nothave explicit overrides for a type will raise aTypeError when an instanceof such a type is passed:

>>>torch.mul(s,3)TypeError: no implementation found for 'torch.mul' on types thatimplement __torch_function__: [ScalarTensor]

In practice this means that if you would like to implement your overrides usinga__torch_function__ implementation along these lines, you will need toexplicitly implement the fulltorch API or the entire subset of the APIthat you care about for your use case. This may be a tall order as the fulltorch API is quite extensive.

Another option is to not returnNotImplemented for operations that are nothandled but to instead pass aTensor to the originaltorchfunction when no override is available. For example, if we change ourimplementation of__torch_function__ forScalarTensor to the one below:

@classmethoddef__torch_function__(cls,func,types,args=(),kwargs=None):ifkwargsisNone:kwargs={}iffuncnotinHANDLED_FUNCTIONSornotall(issubclass(t,(torch.Tensor,ScalarTensor))fortintypes):args=[a.tensor()ifhasattr(a,'tensor')elseaforainargs]returnfunc(*args,**kwargs)returnHANDLED_FUNCTIONS[func](*args,**kwargs)

Thentorch.mul() will work correctly, although the return type will alwaysbe aTensor rather than aScalarTensor, even if both operandsareScalarTensor instances:

>>>s=ScalarTensor(2,2)>>>torch.mul(s,s)tensor([[4., 0.],        [0., 4.]])

Also see theMetadataTensor example below for another variation on thispattern but instead always returns aMetadataTensor to propagate metadatathrough operations in thetorch API.

The__torch_function__ protocol is designed for full coverage of the API,partial coverage may lead to undesirable results, in particular, certainfunctions raising aTypeError. This is especially true for subclasses,where all three oftorch.add,torch.Tensor.__add__ andtorch.Tensor.addmust be covered, even if they return exactly the same result. Failing to dothis may also lead to infinite recursion. If one requires the implementationof a function fromtorch.Tensor subclasses, they must usesuper().__torch_function__ inside their implementation.

Subclassingtorch.Tensor#

As of version 1.7.0, methods ontorch.Tensor and functions in publictorch.* namespaces applied ontorch.Tensor subclasseswill return subclass instances instead oftorch.Tensor instances:

>>>classSubTensor(torch.Tensor):...pass>>>type(torch.add(SubTensor([0]),SubTensor([1]))).__name__'SubTensor'>>>type(torch.add(SubTensor([0]),torch.tensor([1]))).__name__'SubTensor'

If multiple subclasses exist, the lowest one in the hierarchy will be chosen bydefault. If there is no unique way to determine such a case, then aTypeError is raised:

>>>type(torch.add(SubTensor2([0]),SubTensor([1]))).__name__'SubTensor2'>>>type(torch.add(SubTensor2([0]),torch.tensor([1]))).__name__'SubTensor2'>>>torch.add(SubTensor([0]),OtherSubTensor([1]))Traceback (most recent call last):  File"<stdin>", line1, in<module>TypeError:no implementation found for 'torch.add' on types that implement __torch_function__: [SubTensor, OtherSubTensor]

If one wishes to have a global override for all tensor methods, one can use__torch_function__. Here is an example that logs all function/methodcalls:

classLoggingTensor(torch.Tensor):@classmethoddef__torch_function__(cls,func,types,args=(),kwargs=None):# NOTE: Logging calls Tensor.__repr__, so we can't log __repr__ without infinite recursioniffuncisnottorch.Tensor.__repr__:logging.info(f"func:{func.__name__}, args:{args!r}, kwargs:{kwargs!r}")ifkwargsisNone:kwargs={}returnsuper().__torch_function__(func,types,args,kwargs)

However, if one instead wishes to override a method on the Tensor subclass,there one can do so either by directly overriding the method (by definingit for a subclass), or by using__torch_function__ and matching withfunc.

One should be careful within__torch_function__ for subclasses to alwayscallsuper().__torch_function__(func,...) instead offunc directly,as was the case before version 1.7.0. Failing to do this may causefuncto recurse back into__torch_function__ and therefore cause infiniterecursion.

Extendingtorch with aTensor wrapper type#

Another useful case is a type that wraps aTensor, either as anattribute or via subclassing. Below we implement a special case of this sort oftype, aMetadataTensor that attaches a dictionary of metadata to aTensor that is propagated throughtorch operations. Since thisis a generic sort of wrapping for the fulltorch API, we do not need toindividually implement each override so we can make the__torch_function__implementation more permissive about what operations are allowed:

classMetadataTensor(object):def__init__(self,data,metadata=None,**kwargs):self._t=torch.as_tensor(data,**kwargs)self._metadata=metadatadef__repr__(self):return"Metadata:\n{}\n\ndata:\n{}".format(self._metadata,self._t)@classmethoddef__torch_function__(cls,func,types,args=(),kwargs=None):ifkwargsisNone:kwargs={}metadatas=tuple(a._metadataforainargsifhasattr(a,'_metadata'))args=[getattr(a,'_t',a)forainargs]assertlen(metadatas)>0ret=func(*args,**kwargs)returnMetadataTensor(ret,metadata=metadatas[0])

This simple implementation won’t necessarily work with every function in thetorch API but it is good enough to capture most common operations:

>>>metadata={'owner':'Ministry of Silly Walks'}>>>m=MetadataTensor([[1,2],[3,4]],metadata=metadata)>>>t=torch.tensor([[1,2],[1,2]])>>>torch.add(t,m)Metadata:{'owner': 'Ministry of Silly Walks'}data:tensor([[2, 4],        [4, 6]])>>>torch.mul(t,m)Metadata:{'owner': 'Ministry of Silly Walks'}data:tensor([[1, 4],        [3, 8]])

Operations on multiple types that define__torch_function__#

It is possible to use the torch API with multiple distinct types that each havea__torch_function__ implementation, but special care must be taken. In sucha case the rules are:

• The dispatch operation gathers all distinct implementations of__torch_function__ for each operand and calls them in order: subclassesbefore superclasses, and otherwise left to right in the operator expression.

• If any value other thanNotImplemented is returned, that value isreturned as the result. Implementations can register that they do notimplement an operation by returningNotImplemented.

• If all of the__torch_function__ implementations returnNotImplemented, PyTorch raises aTypeError.

Testing Coverage of Overrides for the PyTorch API#

One troublesome aspect of implementing__torch_function__ is that if someoperations do and others do not have overrides, users will at best see aninconsistent experience, or at worst will see errors raised at runtime when theyuse a function that does not have an override. To ease this process, PyTorchprovides a developer-facing API for ensuring full support for__torch_function__ overrides. This API is private and may be subject tochanges without warning in the future.

First, to get a listing of all overridable functions, usetorch.overrides._get_overridable_functions. This returns a dictionary whosekeys are namespaces in thePyTorch Python API and whose values are a list offunctions in that namespace that can be overridden. For example, let’s print thenames of the first 5 functions intorch.nn.functional that can beoverridden:

>>>fromtorch.overridesimportget_overridable_functions>>>func_dict=get_overridable_functions()>>>nn_funcs=func_dict[torch.nn.functional]>>>print([f.__name__forfinnn_funcs[:5])['adaptive_avg_pool1d', 'adaptive_avg_pool2d', 'adaptive_avg_pool3d', 'adaptive_max_pool1d', 'adaptive_max_pool1d_with_indices']

This listing of functions makes it possible to iterate over all overridablefunctions, however in practice this is not enough to write tests for all ofthese functions without laboriously and manually copying the signature of eachfunction for each test. To ease this process, thetorch.overrides._get_testing_overrides function returns a dictionary mappingoverridable functions in thePyTorch API to dummy lambda functions that havethe same signature as the original function but unconditionally return -1. Thesefunctions are most useful to use withinspect to analyze the functionsignature of the originalPyTorch function:

>>>importinspect>>>fromtorch.overridesimportget_testing_overrides>>>override_dict=get_testing_overrides()>>>dummy_add=override_dict[torch.add]>>>inspect.signature(dummy_add)<Signature (input, other, out=None)>

Finally,torch.overrides.get_ignored_functions returns a tuple of functionsthat explicitly cannot be overridden by__torch_function__. This list can beuseful to confirm that a function that isn’t present in the dictionary returnedbyget_overridable_functions cannot be overridden.

__torch_dispatch__ calling convention#

@classmethoddef__torch_dispatch__(cls,func,types,args=(),kwargs=None):pass

When a user calls an operator with inputs that have__torch_dispatch__, that callmay be forwarded to the__torch_dispatch__. args and kwargs get normalized beforethe call to__torch_dispatch__, that is:

• thekwargs consist of keyword-only arguments in the operator’s schema.If a kwarg is equal to its default value (in the schema), it will not be passed.

• theargs consists of all other arguments, no matter how they were passedto the operator (positional vs keyword).If an arg is equal to its default value, andit is the right-most positional arg or all the args to the right of itare not passed, it will not be passed.