compiled() — pytorch Function Reference

Source Code

benchmarks/dynamo/genai_layers/kernels.py lines 494–503

    def compiled(self, args, kwargs=None) -> Any:
        if kwargs is not None:
            raise AssertionError(f"Expected kwargs to be None, but got {kwargs}")
        x, w, dy = args
        y = torch.compile(self.rms_norm_ref, mode=self.compile_mode, fullgraph=True)(
            x, w
        )
        return lambda: torch.autograd.grad(
            y, [x, w], grad_outputs=dy, retain_graph=True
        )