def METHOD_NAME(self):
    torch.utils.rename_privateuse1_backend("extension_device")
    register_backend_for_device(
        "extension_device", ExtensionScheduling, ExtensionWrapperCodegen
    )
    self.assertTrue(
        get_scheduling_for_device("extension_device") == ExtensionScheduling
    )
    self.assertTrue(
        get_wrapper_codegen_for_device("extension_device")
        == ExtensionWrapperCodegen
    )
    self.assertFalse(self.module.custom_op_called())
    device = self.module.custom_device()
    x = torch.empty(2, 16).to(device=device).fill_(1)
    self.assertTrue(self.module.custom_op_called())
    y = torch.empty(2, 16).to(device=device).fill_(2)
    z = torch.empty(2, 16).to(device=device).fill_(3)
    ref = torch.empty(2, 16).fill_(5)
    self.assertTrue(x.device == device)
    self.assertTrue(y.device == device)
    self.assertTrue(z.device == device)
    def fn(a, b, c):
        return a * b + c
    metrics.reset()
    opt_fn = torch.compile()(fn)
    code = run_and_get_cpp_code(opt_fn, x, y, z)
    FileCheck().check("void kernel").check("loadu").check("extension_device").run(
        code
    )
    opt_fn(x, y, z)
    res = opt_fn(x, y, z)
    self.assertEqual(ref, res.to(device="cpu"))