[sharktank] Fix module patching and safetensors comparison tool

sharktank/sharktank/tools/compare_safetensors.py

@@ -145,7 +145,8 @@ def main(argv):
         type=Path,
         default=None,
         help=(
-            "Path to JSON or YAML file that maps names of expected tensors to names of actual tensors."
+            "Path to JSON or YAML file that maps names of expected tensors to names of "
+            'actual tensors. E.g. {"a": "b", "c": "d"}'
         ),
     )
     parser.add_argument(
@@ -179,7 +180,9 @@ def main(argv):
                     act_name = expected_actual_to_key_map[name]
                 if act_name not in act_keys:
                     continue
-            reporter.compare(name, exp_f.get_tensor(name), act_f.get_tensor(act_name))
+                reporter.compare(
+                    name, exp_f.get_tensor(name), act_f.get_tensor(act_name)
+                )
     finally:
         reporter.close()
 

sharktank/sharktank/utils/patching.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Any
 from collections.abc import Mapping, Iterable
+from sharktank.types import InferenceTensor, unbox_tensor
 import logging
 import torch
 
@@ -124,6 +125,7 @@ def save_file(self, output_path: Path, *, skip_unsupported_dtypes: bool = False)
                     for k, v in tensor_dict.items()
                     if k not in unsupported_tensor_dict.keys()
                 }
+
         save_file(tensor_dict, output_path)
 
     def _add_nested_tensors(
@@ -132,16 +134,16 @@ def _add_nested_tensors(
         tensors: list[Any] | dict[str, Any] | torch.Tensor,
         name_delimiter: str,
     ):
-        if isinstance(tensors, torch.Tensor):
-            self._add_tensor(name=name_prefix, tensor=tensors)
+        if isinstance(tensors, (torch.Tensor, InferenceTensor)):
+            self._add_tensor(name=name_prefix, tensor=unbox_tensor(tensors))
         elif isinstance(tensors, Mapping):
             for k, v in tensors.items():
-                self._add_nested_tensor(
+                self._add_nested_tensors(
                     f"{name_prefix}{name_delimiter}{k}", v, name_delimiter
                 )
         elif isinstance(tensors, Iterable):
             for i, v in enumerate(tensors):
-                self._add_nested_tensor(
+                self._add_nested_tensors(
                     f"{name_prefix}{name_delimiter}{i}", v, name_delimiter
                 )
         else:
@@ -154,7 +156,6 @@ def _add_tensor(self, name: str, tensor: torch.Tensor):
             del self.tensors[name]
             self.duplicate_tensors[name] = 0
             self.tensors[f"{name}#0"] = orig_dup
-
         if name in self.duplicate_tensors:
             index = self.duplicate_tensors[name] + 1
             self.duplicate_tensors[name] = index