Fixes for:

b/498080408
b/498074854
b/498090408
b/500231278

Author: QuinnMMcGarry

axlearn/common/attention_test.py

@@ -3479,8 +3479,8 @@ def test_scale_query_key_barrier(self):
         self.assertIn(str(key_scale_factor), hlo)
         # This test is unsupported on Jax 0.8.2 and later
         # TODO(rohit-chatterjee): Rewrite or disable test entirely
-        if jax.__version__ < "0.8.2":
-            self.assertNotIn(str(query_scale_factor * key_scale_factor), hlo)
+        # if jax.__version__ < "0.8.2":
+        #     self.assertNotIn(str(query_scale_factor * key_scale_factor), hlo)
 
     @parameterized.product(
         [

axlearn/common/flash_attention/tpu_paged_attention.py

@@ -258,8 +258,8 @@ def __call__(
 
         in_specs = [
             q_block_spec,  # Query
-            pl.BlockSpec(memory_space=pltpu.ANY),  # Key pages
-            pl.BlockSpec(memory_space=pltpu.ANY),  # Value pages
+            pl.BlockSpec(memory_space=None),  # Key pages
+            pl.BlockSpec(memory_space=None),  # Value pages
             bias_spec,  # Bias
             logit_sink_spec,  # Logit sink
         ]

axlearn/common/pipeline.py

@@ -288,7 +288,7 @@ def scan(
         """See `BaseSchedule.scan` for details."""
 
         @functools.partial(
-            jax.ad_checkpoint.checkpoint,
+            jax.checkpoint,
             prevent_cse=False,
             policy=jax.checkpoint_policies.nothing_saveable,
         )
@@ -515,7 +515,7 @@ def scan(
         m = self.num_microbatches
 
         @functools.partial(
-            jax.ad_checkpoint.checkpoint,
+            jax.checkpoint,
             prevent_cse=False,
             policy=jax.checkpoint_policies.save_only_these_names("iter_input"),
         )

axlearn/common/rnn.py

@@ -208,7 +208,7 @@ def extend_step(
         old_c = cached_states["c"]
         new_c = jax.nn.sigmoid(proj_f) * old_c + jax.nn.sigmoid(proj_g) * jax.nn.tanh(proj_i)
         if cfg.max_cell_value is not None:
-            new_c = jnp.clip(new_c, a_min=-cfg.max_cell_value, a_max=cfg.max_cell_value)
+            new_c = jnp.clip(new_c, min=-cfg.max_cell_value, max=cfg.max_cell_value)
         # [batch, output_dim].
         # Output_nonlinearity is default to True in Lingvo.
         # https://github.com/tensorflow/lingvo/blob/06553f297bbc38eb369503a421d07515d114cbb4/lingvo/core/rnn_cell.py#L247.

axlearn/common/ssm.py

@@ -1519,7 +1519,7 @@ def initialize(
         ).astype(
             dtype
         )  # math.log may return float64, so we need to cast to dtype
-        dt = jnp.clip(dt, a_min=cfg.dt_init_floor)
+        dt = jnp.clip(dt, min=cfg.dt_init_floor)
         # Get inverse of softplus.
         inv_dt = dt + jnp.log(-jnp.expm1(-dt))
         return inv_dt

axlearn/common/trainer_test.py

@@ -535,11 +535,11 @@ def mock_compile_train_step(*args, compiler_options=None, **kwargs):
                 if not enable_python_cache:
                     # As of Jax >= 0.6.0, enable_python_cache=False no longer affects the
                     # AOT compilation path. We now expect the cache hits to be 0
-                    if jax.__version__ >= "0.6.0":
-                        pytest.skip(
-                            # pylint: disable-next=line-too-long
-                            "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
-                        )
+                    # if jax.__version__ >= "0.6.0":
+                    pytest.skip(
+                        # pylint: disable-next=line-too-long
+                        "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
+                    )
                     # We expect to have hit the lowering cache on all but one step.
                     self.assertEqual(end_cache_hits - start_cache_hits, cfg.max_step - 1)
                     self.assertEqual(mocked_compile_fn.call_count, cfg.max_step)
@@ -551,11 +551,11 @@ def mock_compile_train_step(*args, compiler_options=None, **kwargs):
                 self.assertEqual(compiled_with_options_call_count[0], 2)
             else:
                 if not enable_python_cache:
-                    if jax.__version__ >= "0.6.0":
-                        pytest.skip(
-                            # pylint: disable-next=line-too-long
-                            "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
-                        )
+                    # if jax.__version__ >= "0.6.0":
+                    pytest.skip(
+                        # pylint: disable-next=line-too-long
+                        "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
+                    )
                     self.assertEqual(end_cache_hits - start_cache_hits, cfg.max_step - 1)
                     self.assertEqual(mocked_compile_fn.call_count, cfg.max_step)
                 else:

axlearn/huggingface/hf_extractive_qa.py

@@ -277,8 +277,8 @@ def forward(  # pylint: disable=duplicate-code
         if start_positions is not None and end_positions is not None:
             # If the start/end positions are outside of model inputs, we ignore these terms.
             seq_len = start_logits.shape[1]
-            start_positions = jnp.clip(start_positions, a_max=seq_len - 1)
-            end_positions = jnp.clip(end_positions, a_max=seq_len - 1)
+            start_positions = jnp.clip(start_positions, max=seq_len - 1)
+            end_positions = jnp.clip(end_positions, max=seq_len - 1)
 
             is_valid_input = jnp.logical_and(start_positions >= 0, end_positions >= 0)
             num_inputs = is_valid_input.sum()

axlearn/vision/clip.py

@@ -522,7 +522,7 @@ def forward(self, input_batch: NestedTensor) -> NestedTensor:
         log_logit_scale = self.parameters["log_logit_scale"]
         # Ref: https://arxiv.org/pdf/2103.00020.pdf pp.5 Sect. 2.5
         # Clip to prevent scaling the logits by more than 100.
-        log_logit_scale = jnp.clip(log_logit_scale, a_max=jnp.log(cfg.temperature_max_cap))
+        log_logit_scale = jnp.clip(log_logit_scale, max=jnp.log(cfg.temperature_max_cap))
         temperature = 1 / jnp.exp(log_logit_scale)
         similarity = contrastive_logits(x, y)
         loss = self._contrastive_loss_fn(similarity, similarity.T, temperature=temperature)