get around of +/- inf issue when compile flatbuffer file (#2460)

Summary: Pull Request resolved: #2460

Reviewed By: guangy10

Differential Revision: D54961333

fbshipit-source-id: 31fc556aafa0421e7537f19dc2a367739f510003

Author: Gasoonjia

exir/_serialize/_flatbuffer.py

@@ -11,12 +11,12 @@
 import re
 import shutil
 import subprocess
+
 import tempfile
 
 from dataclasses import dataclass
 from typing import Callable, Dict, List, Optional, Sequence
 
-
 # If this environment variable is set to true, save the flatc input files when
 # serialization fails.
 _SAVE_FLATC_ENV: str = "ET_EXIR_SAVE_FLATC_INPUTS_ON_FAILURE"
@@ -29,6 +29,14 @@ 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],
@@ -281,8 +289,11 @@ 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(program_json.encode("ascii"))
+            json_file.write(replaced_program_json.encode("ascii"))
 
         try:
             _flatc_compile(temp_dir, schema_info.root_path, json_path)