fix fake export

Author: xadupre

onnx_diagnostic/export/dynamic_shapes.py

@@ -226,7 +226,7 @@ def _valid_shapes_tensor(cls, inputs, ds):
         for i, d in enumerate(inputs.shape):
             if i in ds and not isinstance(ds[i], int):
                 # dynamic then
-                if d in {0, 1}:
+                if isinstance(d, int) and d in {0, 1}:
                     # export issues for sure
                     issues[i] = f"d=[{d}]"
         return issues if issues else None

onnx_diagnostic/export/shape_helper.py

@@ -306,6 +306,8 @@ def make_fake_with_dynamic_dimensions(
         return x, fake_mode
     if hasattr(x, "shape"):
         t = fake_reshape(x, dynamic_shapes, fake_mode=fake_mode)
+        assert t.device == x.device, f"device mismatch {x.device} -> {t.device}"
+        assert t.dtype == x.dtype, f"dtype mismatch {x.dtype} -> {t.dtype}"
         return t, fake_mode
     from ..helpers import string_type
 

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -34,7 +34,9 @@ def fake_reshape(
         if true_tensor.shape[i] <= 1:
             expanded_shape = list(true_tensor.shape)
             expanded_shape[i] = _unique()
-            true_tensor = torch.empty(tuple(expanded_shape), dtype=true_tensor.dtype)
+            true_tensor = torch.empty(
+                tuple(expanded_shape), dtype=true_tensor.dtype, device=true_tensor.device
+            )
 
     # deal with equivalent dimension
     new_shape = list(true_tensor.shape)
@@ -47,7 +49,9 @@ def fake_reshape(
             d = _unique()
             mapping[d] = s
             new_shape[i] = d
-    true_tensor = torch.empty(tuple(new_shape), dtype=true_tensor.dtype)
+    true_tensor = torch.empty(
+        tuple(new_shape), dtype=true_tensor.dtype, device=true_tensor.device
+    )
 
     # now switch to FakeTensor
     if fake_mode is None:

onnx_diagnostic/torch_models/validate.py

@@ -1909,6 +1909,14 @@ def call_torch_export_custom(
     strict = "-strict" in exporter
     args, kwargs = split_args_kwargs(data["inputs_export"])
     ds = data.get("dynamic_shapes", None)
+    if "-fake" in exporter:
+        from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
+
+        if verbose:
+            print("[call_torch_export_custom] switching to FakeTensor")
+        assert not args, f"Exporter {exporter!r} not implemented with fake tensors."
+        kwargs = torch_deepcopy(kwargs)
+        kwargs, _ = make_fake_with_dynamic_dimensions(kwargs, dynamic_shapes=ds)
     opset = data.get("model_opset", None)
     if opset:
         summary["export_opset"] = opset