Update schema to include infinity for double values (#5203)

Summary:
Pull Request resolved: #5203

Replace previous solution of parsing the json and replacing instances of "Infinity" with "inf"

Differential Revision: D62393242

Author: lucylq

exir/_serialize/_dataclass.py

@@ -95,6 +95,7 @@ class Example
     cls_flds = fields(cls)
     data = {}
     for field in cls_flds:
+        print("field: ", field.name, "type: ", field.type, "default: ", field.default)
         key = field.name
         T = field.type
 
@@ -129,6 +130,11 @@ class Example
             data[key] = [_json_to_dataclass(e, T) for e in value]
             continue
 
+        if get_origin(T) is Union:
+            res = [x for x in get_args(get_type_hints(cls)[key]) if x == type(value)]
+            data[key] = res[0](value)
+            continue
+
         # If T is an enum then lookup the value in the enum otherwise try to
         # cast value to whatever type is required
         if isinstance(T, enum.EnumMeta):

exir/_serialize/_flatbuffer.py

@@ -29,14 +29,6 @@ def _is_valid_alignment(alignment: int) -> bool:
     return alignment > 0 and (alignment & (alignment - 1)) == 0
 
 
-# TODO(T182299196): Replace this hack with a proper flatc binary.
-def _replace_infinity_in_json_file(content: str) -> str:
-    content = re.sub(
-        r'"double_val"\s*:\s*(-)?Infinity', r'"double_val": "\g<1>inf"', content
-    )
-    return content
-
-
 def _patch_schema_alignment(
     schema: bytes,
     constant_tensor_alignment: Optional[int],
@@ -291,11 +283,8 @@ def _program_json_to_flatbuffer(
         json_path = os.path.join(temp_dir, file_stem + ".json")
         output_path = os.path.join(temp_dir, file_stem + ".pte")
 
-        # TODO(T182299196): Replace this hack with a proper flatc binary.
-        replaced_program_json = _replace_infinity_in_json_file(program_json)
-
         with open(json_path, "wb") as json_file:
-            json_file.write(replaced_program_json.encode("ascii"))
+            json_file.write(program_json.encode("ascii"))
 
         try:
             _flatc_compile(temp_dir, schema_info.root_path, json_path)
@@ -330,6 +319,13 @@ def _program_json_to_flatbuffer(
             )
 
 
+def _replace_infinity_in_json_file(content: str) -> str:
+    content = re.sub(
+        r'"double_val"\s*:\s*(-)?inf', r'"double_val": "\g<1>inf"', content
+    )
+    return content
+
+
 def _program_flatbuffer_to_json(program_flatbuffer: bytes) -> bytes:
     """Converts binary flatbuffer data into Program-compatible JSON.
 
@@ -347,5 +343,7 @@ def _program_flatbuffer_to_json(program_flatbuffer: bytes) -> bytes:
             bin_file.write(program_flatbuffer)
 
         _flatc_decompile(temp_dir, schema_info.root_path, bin_path)
-        with open(json_path, "rb") as output_file:
-            return output_file.read()
+        with open(json_path, "r") as output_file:
+            json_data = output_file.read()
+            replaced_json_data = _replace_infinity_in_json_file(json_data)
+            return bytes(replaced_json_data, "ascii")

exir/_serialize/_program.py

@@ -11,7 +11,7 @@
 import re
 
 from dataclasses import dataclass
-from typing import ClassVar, List, Literal, Optional, Tuple
+from typing import ClassVar, List, Literal, Optional, Tuple, Union
 
 from executorch.exir._serialize._cord import Cord
 from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
@@ -47,7 +47,12 @@ def _program_to_json(program: Program) -> str:
 def _json_to_program(program_json: bytes) -> Program:
     """Returns a Program deserialized from the given JSON string."""
     # construct program class recursively from dict
-    return _json_to_dataclass(json.loads(program_json), cls=Program)
+    print("program_json")
+    print(program_json)
+    x = json.loads(program_json)
+    print("json.loads(program_json)")
+    print(x)
+    return _json_to_dataclass(x, cls=Program)
 
 
 def _padding_required(offset: int, alignment: int) -> int:
@@ -571,9 +576,8 @@ def deserialize_pte_binary(program_data: bytes) -> Program:
         segment_base_offset = eh.segment_base_offset
 
     # Parse the flatbuffer data.
-    program: Program = _json_to_program(
-        _program_flatbuffer_to_json(program_data[:program_size])
-    )
+    x = _program_flatbuffer_to_json(program_data[:program_size])
+    program: Program = _json_to_program(x)
 
     if segment_base_offset != 0:
         # Move segment data back into the Program.

exir/emit/test/test_emit.py

@@ -23,6 +23,7 @@
     ExecutorchProgramManager,
     to_edge,
 )
+from executorch.exir._serialize._program import deserialize_pte_binary
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -35,6 +36,7 @@
 from executorch.exir.schema import (
     Bool,
     DelegateCall,
+    Double,
     EValue,
     ExecutionPlan,
     Int,
@@ -1620,3 +1622,33 @@ def forward(self, x):
         executorch_module = _load_for_executorch_from_buffer(model.buffer)
         self.assertEqual(executorch_module(torch.zeros(1))[0], torch.zeros(1))
         self.assertEqual(executorch_module(torch.zeros(1))[0], torch.zeros(1) + 1)
+
+    def test_infinity_in_model(self) -> None:
+        class InfinityMaskModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.mask = torch.tensor([[1, 0], [0, 1]], dtype=torch.float32)
+
+            def forward(self, x):
+                masked_weights = x.masked_fill(self.mask == 0, float("-inf"))
+                return masked_weights
+
+        model = to_edge(
+            export(
+                InfinityMaskModel(),
+                (torch.randn(2, 2),),
+            )
+        )
+
+        # Confirm that we can serialize the model with infinity in it.
+        model = model.to_executorch()
+
+        # Assert that the infinity is stored as a string "-inf".
+        values = model.executorch_program.execution_plan[0].values
+        self.assertEqual(values[5].val, Double(double_val=float("-inf")))
+
+        # Confirm that we can also deserialize the model with infinity in it.
+        pte_data = deserialize_pte_binary(model.buffer)
+        self.assertEqual(
+            pte_data.execution_plan, model.executorch_program.execution_plan
+        )

exir/schema.py

@@ -75,7 +75,23 @@ class Bool:
 
 @dataclass
 class Double:
-    double_val: float
+    double_val: Union[float, str]
+
+    def __init__(self, double_val: float) -> None:
+        if double_val == float("inf"):
+            self.double_val = "inf"
+        elif double_val == float("-inf"):
+            self.double_val = "-inf"
+        else:
+            self.double_val = double_val
+
+    def __post_init__(self) -> None:
+        if isinstance(self.double_val, str):
+            assert self.double_val in ["inf", "-inf"]
+        else:
+            assert isinstance(self.double_val, float)
+            assert not self.double_val == float("inf")
+            assert not self.double_val == float("-inf")
 
 
 @dataclass