[JAX] Replace jnp.clip(..., a_min=..., a_max=...) with jnp.clip(..., min=..., max=...).

a_min and a_max are deprecated parameter names to jax.numpy.clip.

PiperOrigin-RevId: 890629221

Author: hawkinsp

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -193,7 +193,9 @@ def prepare_video_coords(
     # pixel_coords[:, 0, ...] selects Frame dimension.
     # pixel_coords shape: [B, 3, num_patches, 2] -> dim 1 is (F, H, W)
     frame_coords = pixel_coords[:, 0, ...]
-    frame_coords = jnp.clip(frame_coords + self.causal_offset - self.scale_factors[0], a_min=0)
+    frame_coords = jnp.clip(
+        frame_coords + self.causal_offset - self.scale_factors[0], min=0
+    )
     pixel_coords = pixel_coords.at[:, 0, ...].set(frame_coords / fps)
 
     return pixel_coords
@@ -210,12 +212,16 @@ def prepare_audio_coords(
     # 2. Start timestamps
     audio_scale_factor = self.scale_factors[0]
     grid_start_mel = grid_f * audio_scale_factor
-    grid_start_mel = jnp.clip(grid_start_mel + self.causal_offset - audio_scale_factor, a_min=0)
+    grid_start_mel = jnp.clip(
+        grid_start_mel + self.causal_offset - audio_scale_factor, min=0
+    )
     grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate
 
     # 3. End timestamps
     grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor
-    grid_end_mel = jnp.clip(grid_end_mel + self.causal_offset - audio_scale_factor, a_min=0)
+    grid_end_mel = jnp.clip(
+        grid_end_mel + self.causal_offset - audio_scale_factor, min=0
+    )
     grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate
 
     # Stack [num_patches, 2]

src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py

@@ -188,7 +188,7 @@ def preprocess_conditions(
     if mask is not None:
       mask = self.video_processor.preprocess_video(mask, image_size[0], image_size[1])
       mask = jnp.array(np.asarray(mask), dtype=video.dtype)
-      mask = jnp.clip((mask + 1) / 2, a_min=0, a_max=1)
+      mask = jnp.clip((mask + 1) / 2, min=0, max=1)
     else:
       mask = jnp.ones_like(video)