fall back to eager on NotImplementedError (pytorch#107863)

Follow-up to pytorch#107710:

Help  dynamo fall back to eager when compiling unimplemented numpy constructs:

- arrays of strings
- (arg){min, max} for complex types
- various arguments typed as NotImplemented (`np.ones(4, order="F")` etc)
- numpy functions which torch._numpy does not implement

To test, run (we do not implement arrays of strings)

```
import torch
import numpy as np

@torch.compile(fullgraph=False)
def fn():
    return np.asarray(["L", "U"])
```

and observe it compiles with fullgraph=False and fails with fullgraph=True

Fixes pytorch#107970

Pull Request resolved: pytorch#107863
Approved by: https://github.com/ezyang, https://github.com/lezcano

test/dynamo/test_misc.py

@@ -1574,6 +1574,28 @@ def fn():
         self.assertEqual(r.dtype, np.float16)
         self.assertEqual(cnts.frame_count, 1)
 
+    def test_numpy_fallback_on_eager(self):
+        def fn():
+            return np.asarray(["L", "U"])
+
+        cnts = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch._dynamo.optimize(cnts)(fn)
+
+        r = opt_fn()
+        self.assertEqual(cnts.frame_count, 0)  # graph break
+        self.assertEqual(r, np.asarray(["L", "U"]))
+
+        # repeat with a different function
+        def fn2():
+            return np.random.choice(["L", "U"])
+
+        cnts2 = torch._dynamo.testing.CompileCounter()
+        opt_fn2 = torch._dynamo.optimize(cnts2)(fn2)
+
+        r2 = fn2()
+        self.assertEqual(cnts.frame_count, 0)
+        assert r2 in ("L", "U")
+
     def test_inplace_view_on_graph_input(self):
         # graph break when calling methods with inplace_view tag on graph input
         func_args_map = {

test/test_sparse.py

@@ -591,6 +591,7 @@ def test_sparse_bool(self, device, dtype):
         b = a.to_sparse().to_dense()
         self.assertEqual(a, b)
 
+    @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/108667")
     @dtypes(torch.double, torch.cdouble)
     def test_scalar(self, device, dtype):
         # tensor with value

test/test_torch.py

@@ -4946,6 +4946,7 @@ def compare_strides(s1, s2, div):
 
     @onlyCUDA
     @unittest.skipIf(PYTORCH_CUDA_MEMCHECK, "is_pinned uses failure to detect pointer property")
+    @skipIfTorchDynamo("NotImplementedError: PrimTorch does not support pinned memory")
     def test_pin_memory_from_constructor(self, device):
         def _get_like(t, **kwargs):
             return [

torch/_dynamo/utils.py

@@ -1397,6 +1397,7 @@ def run_node(tracer, node, args, kwargs, nnmodule):
     raise an AssertionError.
     """
     op = node.op
+
     try:
         if op == "call_function":
             return node.target(*args, **kwargs)
@@ -1410,6 +1411,12 @@ def run_node(tracer, node, args, kwargs, nnmodule):
         elif op == "placeholder":
             assert "example_value" in node.meta
             return node.meta["example_value"]
+    except NotImplementedError as e:
+        # NB: mimic how wrap_fake_exception does it
+        from .exc import unimplemented
+
+        raise unimplemented(f"running {op} {node.target}(*{args}, **{kwargs})") from e
+
     except Exception as e:
         fn_str = f"Failed running {op} {node.target}(*{args}, **{kwargs}):\n"
         raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e

torch/_numpy/_funcs_impl.py

@@ -1368,6 +1368,8 @@ def einsum(*operands, out=None, dtype=None, order="K", casting="safe", optimize=
 
 
 def _sort_helper(tensor, axis, kind, order):
+    if tensor.dtype.is_complex:
+        raise NotImplementedError(f"sorting {tensor.dtype} is not supported")
     (tensor,), axis = _util.axis_none_flatten(tensor, axis=axis)
     axis = _util.normalize_axis_index(axis, tensor.ndim)
 
@@ -1377,17 +1379,13 @@ def _sort_helper(tensor, axis, kind, order):
 
 
 def sort(a: ArrayLike, axis=-1, kind=None, order: NotImplementedType = None):
-    if a.dtype.is_complex:
-        return NotImplemented
     # `order` keyword arg is only relevant for structured dtypes; so not supported here.
     a, axis, stable = _sort_helper(a, axis, kind, order)
     result = torch.sort(a, dim=axis, stable=stable)
     return result.values
 
 
 def argsort(a: ArrayLike, axis=-1, kind=None, order: NotImplementedType = None):
-    if a.dtype.is_complex:
-        return NotImplemented
     a, axis, stable = _sort_helper(a, axis, kind, order)
     return torch.argsort(a, dim=axis, stable=stable)
 
@@ -1396,7 +1394,7 @@ def searchsorted(
     a: ArrayLike, v: ArrayLike, side="left", sorter: Optional[ArrayLike] = None
 ):
     if a.dtype.is_complex:
-        return NotImplemented
+        raise NotImplementedError(f"searchsorted with dtype={a.dtype}")
 
     return torch.searchsorted(a, v, side=side, sorter=sorter)
 

torch/_numpy/_normalizations.py

@@ -200,6 +200,7 @@ def wrapped(*args, **kwds):
             sig = inspect.signature(func)
             params = sig.parameters
             first_param = next(iter(params.values()))
+
             # NumPy's API does not have positional args before variadic positional args
             if first_param.kind == inspect.Parameter.VAR_POSITIONAL:
                 args = [maybe_normalize(arg, first_param) for arg in args]
@@ -217,6 +218,7 @@ def wrapped(*args, **kwds):
                 name: maybe_normalize(arg, params[name]) if name in params else arg
                 for name, arg in kwds.items()
             }
+
             result = func(*args, **kwds)
 
             # keepdims

torch/_numpy/_reductions_impl.py

@@ -71,6 +71,9 @@ def argmax(
     *,
     keepdims: KeepDims = False,
 ):
+    if a.is_complex():
+        raise NotImplementedError(f"argmax with dtype={a.dtype}.")
+
     axis = _util.allow_only_single_axis(axis)
 
     if a.dtype == torch.bool:
@@ -88,6 +91,9 @@ def argmin(
     *,
     keepdims: KeepDims = False,
 ):
+    if a.is_complex():
+        raise NotImplementedError(f"argmin with dtype={a.dtype}.")
+
     axis = _util.allow_only_single_axis(axis)
 
     if a.dtype == torch.bool:
@@ -134,6 +140,9 @@ def amax(
     initial: NotImplementedType = None,
     where: NotImplementedType = None,
 ):
+    if a.is_complex():
+        raise NotImplementedError(f"amax with dtype={a.dtype}")
+
     return a.amax(axis)
 
 
@@ -149,6 +158,9 @@ def amin(
     initial: NotImplementedType = None,
     where: NotImplementedType = None,
 ):
+    if a.is_complex():
+        raise NotImplementedError(f"amin with dtype={a.dtype}")
+
     return a.amin(axis)
 
 

torch/_numpy/_util.py

@@ -201,7 +201,10 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
     if isinstance(obj, torch.Tensor):
         tensor = obj
     else:
-        tensor = torch.as_tensor(obj)
+        try:
+            tensor = torch.as_tensor(obj)
+        except Exception:
+            raise NotImplementedError(f"failed to convert {obj} to ndarray")
 
         # tensor.dtype is the pytorch default, typically float32. If obj's elements
         # are not exactly representable in float32, we've lost precision:

torch/_numpy/testing/utils.py

@@ -195,6 +195,10 @@ def assert_equal(actual, desired, err_msg="", verbose=True):
         else:
             return True
 
+    if isinstance(desired, str) and isinstance(actual, str):
+        assert actual == desired
+        return
+
     if isinstance(desired, dict):
         if not isinstance(actual, dict):
             raise AssertionError(repr(type(actual)))
@@ -209,6 +213,7 @@ def assert_equal(actual, desired, err_msg="", verbose=True):
         for k in range(len(desired)):
             assert_equal(actual[k], desired[k], f"item={k!r}\n{err_msg}", verbose)
         return
+
     from torch._numpy import imag, iscomplexobj, isscalar, ndarray, real, signbit
 
     if isinstance(actual, ndarray) or isinstance(desired, ndarray):