Add option to persist_xla_flags from Orbax Native Models.

PiperOrigin-RevId: 833478316

Author: jerryxyj

export/orbax/export/constants.py

@@ -100,6 +100,9 @@ class ExportModelType(enum.Enum):
 # Whether to strip XLA flags from the model.
 STRIP_XLA_FLAGS = 'strip_xla_flags'
 
+# Whether to persist XLA flags in the model.
+PERSIST_XLA_FLAGS = 'persist_xla_flags'
+
 ################################################################################
 # Proto field names
 ################################################################################