more robust

Author: xadupre

_unittests/ut_torch_models/test_test_helpers.py

@@ -15,6 +15,7 @@
     validate_model,
     filter_inputs,
     run_ort_fusion,
+    empty,
 )
 from onnx_diagnostic.tasks import supported_tasks
 
@@ -32,6 +33,9 @@ def test_get_inputs_for_task(self):
             self.assertIn("dynamic_shapes", data)
             copy.deepcopy(data["inputs"])
 
+    def test_empty(self):
+        self.assertFalse(empty("float16"))
+
     @hide_stdout()
     def test_validate_model(self):
         mid = "arnir0/Tiny-LLM"

onnx_diagnostic/helpers/torch_helper.py

@@ -689,9 +689,15 @@ def forward(self, input_ids):
     raise NotImplementedError(f"cls_name={cls_name}")
 
 
-def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
+def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
     """Applies torch.to if applicable. Goes recursively."""
     if isinstance(value, (torch.nn.Module, torch.Tensor)):
+        if (
+            isinstance(to_value, torch.dtype)
+            or to_value in {"float16", "bfloat16", "float32", "float64"}
+        ) and value.dtype in {torch.int32, torch.int64, torch.int8, torch.int16}:
+            # int vector should not be changed.
+            return value
         return value.to(to_value)
     if isinstance(value, list):
         return [to_any(t, to_value) for t in value]
@@ -712,11 +718,20 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
                 )
             )
         )
+    if value.__class__.__name__ == "EncoderDecoderCache":
+        return make_encoder_decoder_cache(
+            to_any(value.self_attention_cache, to_value),
+            to_any(value.cross_attention_cache, to_value),
+        )
     if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
         args, spec = torch.utils._pytree.tree_flatten(value)
         new_args = to_any(args, to_value)
         return torch.utils._pytree.tree_unflatten(new_args, spec)
 
+    assert "Cache" not in value.__class__.__name__, (
+        f"Class {value.__class__.__name__!r} should be registered "
+        f"to be able to change the type in every tensor it contains."
+    )
     assert not isinstance(value, Iterable), f"Unsupported type {type(value)}"
     return value
 

onnx_diagnostic/torch_models/test_helper.py

@@ -23,7 +23,7 @@
 def empty(value: Any) -> bool:
     """Tells if the value is empty."""
     if isinstance(value, (str, list, dict, tuple, set)):
-        return bool(value)
+        return not bool(value)
     if value is None:
         return True
     return False