make etrecord set representive IO

representive input and reference output in etrecord will not be set during export flow. To continue supporting the two functionalities, this diff creates two class methods to customize IO in etrecord.

Differential Revision: [D79386896](https://our.internmc.facebook.com/intern/diff/D79386896/)

[ghstack-poisoned]

Author: Gasoonjia

devtools/etrecord/_etrecord.py

@@ -68,7 +68,7 @@ def __init__(
             Dict[str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]]]
         ] = None,
         _reference_outputs: Optional[Dict[str, List[ProgramOutput]]] = None,
-        _representative_inputs: Optional[List[ProgramOutput]] = None,
+        _representative_inputs: Optional[List[ProgramInput]] = None,
     ):
         self.exported_program = exported_program
         self.export_graph_id = export_graph_id
@@ -345,6 +345,56 @@ def add_edge_dialect_program(
         # Set the extracted data
         self.edge_dialect_program = processed_edge_dialect_program
 
+    def update_representative_inputs(
+        self,
+        representative_inputs: Union[List[ProgramInput], BundledProgram],
+    ) -> None:
+        """
+        Update the representative inputs in the ETRecord.
+
+        This method allows users to customize the representative inputs that will be
+        included when the ETRecord is saved. The representative inputs can be provided
+        directly as a list or extracted from a BundledProgram.
+
+        Args:
+            representative_inputs: Either a list of ProgramInput objects or a BundledProgram
+                from which representative inputs will be extracted.
+        """
+        if isinstance(representative_inputs, BundledProgram):
+            self._representative_inputs = _get_representative_inputs(
+                representative_inputs
+            )
+        else:
+            self._representative_inputs = representative_inputs
+
+    def update_reference_outputs(
+        self,
+        reference_outputs: Union[
+            Dict[str, List[ProgramOutput]], List[ProgramOutput], BundledProgram
+        ],
+    ) -> None:
+        """
+        Update the reference outputs in the ETRecord.
+
+        This method allows users to customize the reference outputs that will be
+        included when the ETRecord is saved. The reference outputs can be provided
+        directly as a dictionary mapping method names to lists of outputs, as a
+        single list of outputs (which will be treated as {"forward": List[ProgramOutput]}),
+        or extracted from a BundledProgram.
+
+        Args:
+            reference_outputs: Either a dictionary mapping method names to lists of
+                ProgramOutput objects, a single list of ProgramOutput objects (treated
+                as outputs for the "forward" method), or a BundledProgram from which
+                reference outputs will be extracted.
+        """
+        if isinstance(reference_outputs, BundledProgram):
+            self._reference_outputs = _get_reference_outputs(reference_outputs)
+        elif isinstance(reference_outputs, list):
+            self._reference_outputs = {"forward": reference_outputs}
+        else:
+            self._reference_outputs = reference_outputs
+
 
 def _get_reference_outputs(
     bundled_program: BundledProgram,