orbax
orbax copied to clipboard
obm: Use original argument names for TF input signatures.
obm: Use original argument names for TF input signatures.
Previously, the original input names were substituted as "input_0", "input_1", etc.
This change updates the logic to extract input argument names directly from the tf.ConcreteFunction's function_def.signature.input_arg, providing more meaningful names in the OBM representation.
Output names are trickier, as they have different names in FunctionDef and SignatureDef. For now they continue to be generated as "output_0", "output_1", etc.