doc

Author: xadupre

_doc/conf.py

@@ -141,6 +141,7 @@ def linkcode_resolve(domain, info):
     ("py:class", "torch.fx.proxy.TracerBase"),
     ("py:class", "torch.FloatTensor"),
     ("py:class", "torch.LongTensor"),
+    ("py:class", "torch.export._trace.ExportArtifact"),
     ("py:class", "torch.utils._pytree.Context"),
     ("py:class", "torch.utils._pytree.KeyEntry"),
     ("py:class", "torch.utils._pytree.TreeSpec"),

_doc/final/plot_export_tiny_llm_attention_input_observer.py

@@ -122,7 +122,9 @@ def generate_text(
 
 # %%
 # Let's measure the discrepancies.
-data = observer.check_discrepancies(filename, progress_bar=True, atol=1e-2, include_io=True)
+data = observer.check_discrepancies(
+    filename, progress_bar=True, atol=1e-2, include_io=True, skip_none=True
+)
 df = pandas.DataFrame(data)
 df.to_excel("plot_export_tiny_llm_attention_input_observer.xlsx")
 print(df)

_doc/technical/plot_histc.py

@@ -193,7 +193,7 @@ def tune_threshold_histc(
 )
 
 # %%
-# This does not add up. Let's proove now :func:`torch.histc` is really confusing.
+# This does not add up. Let's prove now :func:`torch.histc` is really confusing.
 # The following sum should be null but it is not.
 
 diff = torch.histc(tensor, hbins, hmin, hmax) - (

onnx_diagnostic/helpers/dot_helper.py

@@ -203,6 +203,8 @@ def _mkn(obj: object) -> int:
                 if att.type == onnx.AttributeProto.GRAPH:
                     unique |= get_hidden_inputs(att.g)
             for i in unique:
+                if i in tiny_inits:
+                    continue
                 edge = name_to_ids[i], _mkn(node)  # type: ignore[assignment]
                 if edge in done:
                     continue

onnx_diagnostic/helpers/helper.py

@@ -1071,6 +1071,7 @@ def max_diff(
     _index: int = 0,
     allow_unique_tensor_with_list_of_one_element: bool = True,
     hist: Optional[Union[bool, List[float]]] = None,
+    skip_none: bool = False,
 ) -> Dict[str, Union[float, int, Tuple[Any, ...]]]:
     """
     Returns the maximum discrepancy.
@@ -1087,6 +1088,7 @@ def max_diff(
     :param allow_unique_tensor_with_list_of_one_element:
         allow a comparison between a single tensor and a list of one tensor
     :param hist: compute an histogram of the discrepancies
+    :param skip_none: skips none value
     :return: dictionary with many values
 
     * abs: max absolute error
@@ -1112,6 +1114,7 @@ def max_diff(
         end=end,
         _index=_index,
         hist=hist,
+        skip_none=skip_none,
     )
     _dkws = {**_dkws_, "flatten": flatten}
     _dkwsf = {**_dkws_, "flatten": False}
@@ -1129,6 +1132,7 @@ def max_diff(
                 debug_info=debug_info,
                 allow_unique_tensor_with_list_of_one_element=False,
                 hist=hist,
+                skip_none=skip_none,
             )
         return max_diff(
             expected,
@@ -1142,6 +1146,7 @@ def max_diff(
             _index=_index,
             allow_unique_tensor_with_list_of_one_element=False,
             hist=hist,
+            skip_none=skip_none,
         )
 
     if expected.__class__.__name__ == "CausalLMOutputWithPast":
@@ -1269,6 +1274,7 @@ def max_diff(
                 _index=_index + ip,
                 flatten=flatten,
                 hist=hist,
+                skip_none=skip_none,
             )
             am = max(am, d["abs"])
             dn = max(dn, d["dnan"])
@@ -1793,6 +1799,9 @@ def max_diff(
             **_dkws,
         )
 
+    if skip_none and (expected is None or got is None):
+        return {"abs": 0, "rel": 0, "dnan": 0, "n": 0, "sum": 0}
+
     raise AssertionError(
         f"Not implemented with implemented with expected="
         f"{string_type(expected)} ({type(expected)}), got={string_type(got)},\n"

onnx_diagnostic/investigate/input_observer.py

@@ -286,22 +286,28 @@ class InputObserverInfo:
     and the arguments to send to :func:`torch.export.export`.
 
     Args:
-        signature_names: Names of the arguments of the method
+        signature_names:
+            Names of the arguments of the method
             the collector tensors come from. They are used if it becomes
             necessary to move positional arguments to named ones.
             They are used a second time because :func:`torch.export.export`
             cares about the order in kwargs and dynamic shapes, it needs
             to be the same in the ordered dictionaries `add_inputs` receive.
-        default_values: Default values defined by the signature of the function,
+        default_values:
+            Default values defined by the signature of the function,
             any value equal to that is ignore to simplify the export.
-        missing: If a named argument (in kwargs) is missing,
+        missing:
+            If a named argument (in kwargs) is missing,
             a default value will be taken in this dictionary,
             this is used when after the prefill step, an argument
             disappears (such as `pixel_values`) and another one
             is added (such as `past_key_values`).
             The values are only to infer dynamic shapes and arguments,
             not to run the model.
-        kwargs_name: Name of parameter **kwargs if it exists.
+        kwargs_name:
+            Name of parameter `**kwargs` if it exists.
+
+    This is used by class :class:`InputObserver`.
     """
 
     def __init__(
@@ -910,6 +916,7 @@ def check_discrepancies(
         hist=(0.1, 0.01),
         progress_bar: bool = False,
         include_io: bool = True,
+        skip_none: bool = True,
     ) -> list[dict[str, str | int | float | bool]]:
         """Computes the discrepancies between the saved inputs and outputs
         with the saved onnx model.
@@ -929,6 +936,8 @@ def check_discrepancies(
             include_io:
                 Shows inputs/outputs shapes in the summary
                 returned by this function.
+            skip_none:
+                Dooes not check discrepancies when an output is None.
 
         Returns:
             A list of dictionaries, ready to be consumed by a dataframe.
@@ -982,7 +991,7 @@ def check_discrepancies(
                 if isinstance(outputs, list) and isinstance(ort_outputs, list):
                     while len(ort_outputs) > len(outputs) and ort_outputs[-1].numel() == 0:
                         ort_outputs.pop()
-                diff = max_diff(outputs, ort_outputs, hist=lhist)  # type: ignore[assignment]
+                diff = max_diff(outputs, ort_outputs, hist=lhist, skip_none=skip_none)  # type: ignore[assignment]
                 if "rep" in diff and isinstance(diff["rep"], dict):
                     diff.update(diff["rep"])
                     del diff["rep"]

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -143,7 +143,7 @@ def patched__get_range_constraints(
 ):
     """
     Patches ``torch.export._trace._get_range_constraints``.
-    See PR `#174593 whttps://github.com/pytorch/pytorch/pull/174593>`_.
+    See PR `#174593 <https://github.com/pytorch/pytorch/pull/174593>`_.
     """
     gm: torch.fx.GraphModule = export_artifact.aten.gm
     export_graph_signature: torch.export.graph_signature.ExportGraphSignature = (