From the autograd and amp docs, I gather that the @custom_fwd and @custom_bwd decorators let your custom autograd Function compose nicely w/ mixed precision. In what scenarios, then, would you not wrap your Function’s methods with these decorators?
I’m asking because I was a little surprised that custom Functions and amp don’t compose by default, but I’m sure there’s a good reason for it. Thanks!