Add option to persist_xla_flags from Orbax Native Models.

PiperOrigin-RevId: 833478316

Author: jerryxyj

export/orbax/export/constants.py

@@ -97,9 +97,13 @@ class ExportModelType(enum.Enum):
 # Mesh for the model.
 JAX_MESH = 'jax_mesh'
 
+# TODO: b/459991985 - Remove this flag and use PERSIST_XLA_FLAGS instead.
 # 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
 ################################################################################