Delete the floating point optimization level and memory fitting level settings in favor of the EffortLevel enum.

##### Why?
Simplifies flag input handling by removing multiple ways to specify the same thing.

PiperOrigin-RevId: 821862460

Author: rishipal

jax/_src/compiler.py

@@ -188,8 +188,6 @@ def get_compile_options(
     assert device_assignment.computation_count() == num_partitions
     compile_options.device_assignment = device_assignment
 
-  build_options.exec_time_optimization_effort = config.exec_time_optimization_effort.value
-  build_options.memory_fitting_effort = config.memory_fitting_effort.value
   build_options.optimization_level = config.EffortLevel(
       config.optimization_level.value
   ).value
@@ -199,11 +197,7 @@ def get_compile_options(
 
   if env_options_overrides is not None:
     # Some overrides are passed directly on build_options.
-    overrides_on_build_options = [
-        "exec_time_optimization_effort", "memory_fitting_effort"]
-    overrides_on_build_options.extend(
-        ["optimization_level", "memory_fitting_level"]
-    )
+    overrides_on_build_options = ["optimization_level", "memory_fitting_level"]
 
     env_options_overrides = dict(env_options_overrides)
     for name in overrides_on_build_options:

jax/_src/config.py

@@ -2136,18 +2136,6 @@ def _default_pmap_no_rank_reduction(new_val):
     ),
 )
 
-exec_time_optimization_effort = float_state(
-    name='jax_exec_time_optimization_effort',
-    default=0.0,
-    help='Effort for minimizing execution time (higher means more effort), valid range [-1.0, 1.0].'
-)
-
-memory_fitting_effort = float_state(
-    name='jax_memory_fitting_effort',
-    default=0.0,
-    help='Effort for minimizing memory usage (higher means more effort), valid range [-1.0, 1.0].'
-)
-
 optimization_level = enum_state(
     name='jax_optimization_level',
     enum_values=[

jaxlib/_jax/__init__.pyi

@@ -561,16 +561,6 @@ class ExecutableBuildOptions:
       self, arg: bytes, /
   ) -> None: ...
   @property
-  def exec_time_optimization_effort(self) -> float: ...
-  @exec_time_optimization_effort.setter
-  def exec_time_optimization_effort(
-      self, arg: float, /
-  ) -> ExecutableBuildOptions: ...
-  @property
-  def memory_fitting_effort(self) -> float: ...
-  @memory_fitting_effort.setter
-  def memory_fitting_effort(self, arg: float, /) -> ExecutableBuildOptions: ...
-  @property
   def optimization_level(self) -> int: ...
   @optimization_level.setter
   def optimization_level(self, arg: int, /) -> None: ...

jaxlib/xla_compiler.cc

@@ -1139,12 +1139,6 @@ void BuildXlaCompilerSubmodule(nb::module_& m) {
                  xla::CompilationEnvironments::CreateFromProto(env_proto));
              *options.mutable_comp_envs() = std::move(*comp_envs);
            })
-      .def_prop_rw("exec_time_optimization_effort",
-                   &ExecutableBuildOptions::exec_time_optimization_effort,
-                   &ExecutableBuildOptions::set_exec_time_optimization_effort)
-      .def_prop_rw("memory_fitting_effort",
-                   &ExecutableBuildOptions::memory_fitting_effort,
-                   &ExecutableBuildOptions::set_memory_fitting_effort)
       .def_prop_rw(
           "optimization_level",
           [](ExecutableBuildOptions& options) {

tests/api_test.py

@@ -1383,23 +1383,6 @@ def f(x):
             "xla_gpu_auto_spmd_partitioning_memory_budget_ratio": 0.5,
         })(1.0)  # doesn't crash.
 
-  def test_exec_time_optimization_effort_compiler_option(self):
-    def f(x):
-      return jnp.sqrt(x ** 2) + 1.
-
-    f_jit = jit(
-        f,
-        compiler_options={
-            "exec_time_optimization_effort": 0.0,
-        })(1.0)  # doesn't crash.
-
-    with self.assertRaisesRegex(jax.errors.JaxRuntimeError, "No such"):
-      f_jit = jit(
-          f,
-          compiler_options={
-              "exec_time_compilation_effort": 0.0,
-          })(1.0)
-
   def test_optimization_level_compiler_option(self):
     def f(x):
       return jnp.sqrt(x**2) + 1.0