Move use_colocated_python out of the GOOGLE-INTERNAL block in array_handler_benchmark.yaml.

PiperOrigin-RevId: 834750373

Author: Orbax Authors

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
 ################################################################################

model/orbax/experimental/model/cli/README.md

@@ -2,6 +2,7 @@
 
 A command-line tool for inspecting Orbax models.
 
+
 ## Examples
 
 To inspect the model: