Extra CR comments from #95621 (#96043)

Specifically: https://github.com/pytorch/pytorch/pull/95621/files/063e44147152f4dd7e51852cf8c679692bd9fd53#r1120306196 pytorch/pytorch#95621 (comment) Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch/pytorch#96043 Approved by: https://github.com/Chillee, https://github.com/albanD
cyyever · Mar 12, 2023 · d50ade0 · d50ade0
1 parent 92ef868
commit d50ade0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py
@@ -957,11 +957,10 @@ def same(
             log.error(f"Accuracy failed (float): {ref} != {res} (within tol={tol})")
         return r
     elif is_numpy_int_type(ref) or is_numpy_float_type(ref):
-        if relax_numpy_equality:
-            if is_numpy_int_type(ref):
-                ref = ref.item()
-            if is_numpy_int_type(res):
-                res = res.item()
+        if relax_numpy_equality and not (
+            is_numpy_int_type(res) or is_numpy_float_type(res)
+        ):
+            ref = ref.item()
         r = (type(ref) is type(res)) and (ref == res)
         if not r:
             log.error(f"Accuracy failed (numpy): {ref} != {res}")

diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -787,6 +787,10 @@ def wrap_unspecialized_primitive(self, value):
         if self.name in self.tx.output.unspec_variable_map:
             return self.tx.output.unspec_variable_map[self.name]
         else:
+            # NB: We do not do float.  For motivation, see
+            # https://docs.google.com/document/d/1INSCdYu1PxXcr43HrD82OudeEuS-qxQe1yZmLg2wy6A/edit
+            # but the general idea is that we generate kernels that can
+            # take unspecialized floats and use them in sizevar computation
             if (
                 config.dynamic_shapes
                 and isinstance(value, int)
@@ -799,11 +803,7 @@ def wrap_unspecialized_primitive(self, value):
                 self.tx.output.tracked_fakes.append(
                     TrackedFake(wrapped_value, self.source)
                 )
-                # TODO: Do float?
-                # Not entirely clear we want to do this, as float inputs don't
-                # work with inductor codegen at the moment
             else:
-                # TODO: Eliminate this case entirely
                 wrapped_value = torch.tensor(value)
             if not isinstance(self.get_source(), RandomValueSource):
                 guards = {self.get_source().make_guard(GuardBuilder.TYPE_MATCH, True)}