Update on "[dynamo][not ready] polyfill infra for classes"

anijain2305 · anijain2305 · commit f770084b53e0 · 2025-02-09T15:54:17.000-08:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]
diff --git a/aten/src/ATen/native/cuda/Math.cuh b/aten/src/ATen/native/cuda/Math.cuh
@@ -758,11 +758,10 @@ const auto sinc_string = jiterator_stringify(
   T sinc(T a) {
     if (a == T(0)) {
       return T(1);
-    } else {
-      constexpr T pi = T(3.14159265358979323846L);
-      T product = pi * a;
-      return std::sin(product) / product;
     }
+    constexpr T pi = T(3.14159265358979323846L);
+    T product = pi * a;
+    return std::sin(product) / product;
   }
 ); // sinc_string
 
diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py
@@ -2510,9 +2510,9 @@ def mask_mod(b, h, q, kv):
     @supported_platform
     def test_strided_backwards(self):
         shape = (1, 2, 4096, 64)
-        Q = torch.randn(shape, requires_grad=True, device="cuda", dtype=torch.bfloat16)
-        K = torch.randn(shape, requires_grad=True, device="cuda", dtype=torch.bfloat16)
-        V = torch.randn(shape, requires_grad=True, device="cuda", dtype=torch.bfloat16)
+        Q = torch.randn(shape, requires_grad=True, device="cuda")
+        K = torch.randn(shape, requires_grad=True, device="cuda")
+        V = torch.randn(shape, requires_grad=True, device="cuda")
         func = torch.compile(flex_attention, dynamic=True, fullgraph=True)
 
         K_sliced = K[:, :, :-128]

Original file line number	Diff line number	Diff line change
`@@ -758,11 +758,10 @@ const auto sinc_string = jiterator_stringify(`
`758`	`758`	`T sinc(T a) {`
`759`	`759`	`if (a == T(0)) {`
`760`	`760`	`return T(1);`
`761`		`- } else {`
`762`		`- constexpr T pi = T(3.14159265358979323846L);`
`763`		`- T product = pi * a;`
`764`		`- return std::sin(product) / product;`
`765`	`761`	`}`
	`762`	`+ constexpr T pi = T(3.14159265358979323846L);`
	`763`	`+ T product = pi * a;`
	`764`	`+ return std::sin(product) / product;`
`766`	`765`	`}`
`767`	`766`	`); // sinc_string`
`768`	`767`