remove support for proto formatted xla flags

PiperOrigin-RevId: 811865727

Author: Orbax Authors

model/orbax/experimental/model/core/python/compile_options_util.py

@@ -49,20 +49,16 @@ def generate_tpu_compilation_env(
   )
   # Override with supplied XLA flags if any is provided.
   if xla_flags:
-    is_proto_formatted = False if xla_flags[0].startswith('--') else True
-    if is_proto_formatted:
-      merge_proto_formatted_flags_into_compile_options(xla_flags, env)
-    else:
-      parsed_flags = {}
-      for flag in xla_flags:
-        if not flag.startswith('--'):
-          raise ValueError(
-              f"Flag {flag} does not start with '--'. All flags must be in the"
-              ' format of --flag_name=flag_value when using this format.'
-          )
-        flag_name, flag_value = flag[2:].split('=', 1)
-        parsed_flags[flag_name] = flag_value
-      merge_flags_into_compile_options(parsed_flags, env)
+    parsed_flags = {}
+    for flag in xla_flags:
+      if not flag.startswith('--'):
+        raise ValueError(
+            f"Flag {flag} does not start with '--'. All flags must be in the"
+            ' format of --flag_name=flag_value.'
+        )
+      flag_name, flag_value = flag[2:].split('=', 1)
+      parsed_flags[flag_name] = flag_value
+    merge_flags_into_compile_options(parsed_flags, env)
 
   # Pack the TPU compilation environment into a compilation env proto.
   any_proto = any_pb2.Any()
@@ -231,8 +227,8 @@ def merge_flags_into_compile_options(
   Args:
     xla_flags: A mapping of XLA flag names to their string values. These flags
       will be parsed and merged into the `env` proto.
-    env: The TpuCompilationEnvironment proto to merge the flags into. This
-      proto will be modified in place.
+    env: The TpuCompilationEnvironment proto to merge the flags into. This proto
+      will be modified in place.
   """
   env_override = tpu_comp_env_pb2.TpuCompilationEnvironment()
   for flag_name, value in xla_flags.items():
@@ -245,20 +241,3 @@ def merge_flags_into_compile_options(
       # For scalar types, we can set the attribute directly.
       setattr(env_override, field_descriptor.name, parsed_value)
   env.MergeFrom(env_override)
-
-
-# TODO(b/438187387): remove this path and only allow the "--flag=value" format.
-def merge_proto_formatted_flags_into_compile_options(
-    xla_flags: Sequence[str],
-    env: tpu_comp_env_pb2.TpuCompilationEnvironment,
-):
-  """Merges flags into a proto."""
-  env_override = tpu_comp_env_pb2.TpuCompilationEnvironment()
-  xla_flags_str = '\n'.join(xla_flags)
-  try:
-    text_format.Parse(xla_flags_str, env_override)
-  except text_format.ParseError as e:
-    raise ValueError(
-        f'Error parsing supplied XLA flag overrides {xla_flags_str}.'
-    ) from e
-  env.MergeFrom(env_override)

model/orbax/experimental/model/core/python/compile_options_util_test.py

@@ -41,16 +41,6 @@ def _get_expected_proto_from_tpu_comp_env(field_str: str, proto_str: str):
     'xla_tpu_async_copy_bandwidth_scaling_factor': '0.19125064716453793',
 }
 
-XLA_FLAGS_PROTO_FORMATTED = [
-    'xla_jf_rematerialization_percent_shared_memory_limit: 99',
-    'xla_tpu_allocate_scoped_vmem_at_same_offset: false',
-    (
-        'xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers:'
-        " 'NO_SCALE'"
-    ),
-    'xla_tpu_memory_bound_loop_optimizer_options: {enabled:false}',
-    'xla_tpu_async_copy_bandwidth_scaling_factor: 0.19125064716453793',
-]
 
 EXPECTED_ENV = tpu_comp_env_pb2.TpuCompilationEnvironment(
     xla_jf_rematerialization_percent_shared_memory_limit=99,
@@ -122,19 +112,8 @@ def test_parse_flag_from_string_nonexistent_flag(self):
     with self.assertRaisesRegex(ValueError, 'Flag not found: nonexistent_flag'):
       compile_options_util.parse_flag_from_string('nonexistent_flag', 'value')
 
-  @parameterized.named_parameters(
-      dict(
-          testcase_name='dict_xla_flags',
-          xla_flags=XLA_FLAGS_DICT,
-          merge_fn=compile_options_util.merge_flags_into_compile_options,
-      ),
-      dict(
-          testcase_name='proto_formatted_xla_flags',
-          xla_flags=XLA_FLAGS_PROTO_FORMATTED,
-          merge_fn=compile_options_util.merge_proto_formatted_flags_into_compile_options,
-      ),
-  )
-  def test_merge_flags_into_compile_options(self, xla_flags, merge_fn):
+  def test_merge_flags_into_compile_options(self):
+    xla_flags = XLA_FLAGS_DICT
     # Initialize the environment with some values.
     env = tpu_comp_env_pb2.TpuCompilationEnvironment()
     # Values that should be overridden.
@@ -144,7 +123,7 @@ def test_merge_flags_into_compile_options(self, xla_flags, merge_fn):
     env.xla_tpu_wait_n_cycles_before_program_termination = 1234
 
     # Merge the flags into the environment.
-    merge_fn(xla_flags, env)
+    compile_options_util.merge_flags_into_compile_options(xla_flags, env)
     self.assertEqual(
         env.xla_jf_rematerialization_percent_shared_memory_limit, 99
     )
@@ -170,11 +149,6 @@ def test_merge_flags_into_compile_options(self, xla_flags, merge_fn):
           xla_flags=[f'--{k}={v}' for k, v in XLA_FLAGS_DICT.items()],
           expected_env=EXPECTED_ENV,
       ),
-      dict(
-          testcase_name='proto_formatted_xla_flags',
-          xla_flags=XLA_FLAGS_PROTO_FORMATTED,
-          expected_env=EXPECTED_ENV,
-      ),
       dict(
           testcase_name='no_xla_flags',
           xla_flags=None,
@@ -204,7 +178,7 @@ def test_generate_tpu_compilation_env_invalid_flag_format(self):
         ValueError,
         'Flag xla_tpu_allocate_scoped_vmem_at_same_offset: false does not start'
         " with '--'. All flags must be in the format of"
-        ' --flag_name=flag_value when using this format.',
+        ' --flag_name=flag_value.',
     ):
       compile_options_util.generate_tpu_compilation_env(
           xla_flags=[