Internal change

PiperOrigin-RevId: 941084384

Author: h-joo

recml/core/metrics/confusion_metrics.py

@@ -83,7 +83,7 @@ def compute(self) -> float | Sequence[float]:
     precision_ = _np_divide_no_nan(
         self.true_positives, self.true_positives + self.false_positives
     )
-    return _maybe_squeeze(precision_)
+    return _maybe_squeeze(precision_)  # pyrefly: ignore[bad-return]
 
 
 class Recall(ConfusionMetric):
@@ -93,7 +93,7 @@ def compute(self) -> float | Sequence[float]:
     recall_ = _np_divide_no_nan(
         self.true_positives, self.true_positives + self.false_negatives
     )
-    return _maybe_squeeze(recall_)
+    return _maybe_squeeze(recall_)  # pyrefly: ignore[bad-return]
 
 
 class FBeta(ConfusionMetric):
@@ -145,7 +145,7 @@ def compute(self) -> float | Sequence[float]:
     recall_ = _np_divide_no_nan(
         self.true_positives, self.true_positives + self.false_negatives
     )
-    return _maybe_squeeze(
+    return _maybe_squeeze(  # pyrefly: ignore[bad-return]
         _np_divide_no_nan(
             np.multiply(precision_, recall_) * (self.beta + 1.0),
             np.multiply(precision_, self.beta) + recall_,
@@ -174,7 +174,7 @@ def from_model_output(
         predictions=y_pred,
         labels=y_true,
         weights=weights,
-        thresholds=default_thresholds(num_thresholds),
+        thresholds=default_thresholds(num_thresholds),  # pyrefly: ignore[bad-argument-type]
     )
     return cls(
         true_positives=tp,
@@ -237,7 +237,7 @@ def compute(self) -> float:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # We negate the integral because the thresholds are in ascending order.
-    return -np.trapezoid(tp_rate, fp_rate)
+    return -np.trapezoid(tp_rate, fp_rate)  # pyrefly: ignore[bad-return]
 
 
 class PrecisionAtRecall(ConfusionMetric):
@@ -261,7 +261,7 @@ def from_model_output(
         predictions=y_pred,
         labels=y_true,
         weights=weights,
-        thresholds=default_thresholds(num_thresholds),
+        thresholds=default_thresholds(num_thresholds),  # pyrefly: ignore[bad-argument-type]
     )
     return cls(
         true_positives=tp,
@@ -593,7 +593,7 @@ def _estimate_confusion_matrix(thresholds: jt.Scalar) -> tuple[
 
     return jnp.sum(tp), jnp.sum(tn), jnp.sum(fp), jnp.sum(fn)
 
-  thresholds = jnp.asarray(thresholds, dtype=jnp.float32)
+  thresholds = jnp.asarray(thresholds, dtype=jnp.float32)  # pyrefly: ignore[bad-assignment]
   return jax.vmap(_estimate_confusion_matrix)(thresholds)
 
 

recml/core/metrics/reduction_metrics.py

@@ -39,7 +39,7 @@ def from_fun(cls, fun: Callable[..., Any], **kwargs) -> type[Self]:
     base_cls = cls
     bound_kwargs = kwargs
 
-    class _FromFun(cls):
+    class _FromFun(cls):  # pyrefly: ignore[invalid-inheritance]
       """A reduction metric that is computed from a function."""
 
       @classmethod

recml/core/metrics/tools.py

@@ -112,7 +112,7 @@ def compute_and_log_scalars(
         for k, v in scalars.items()
         if not isinstance(metrics[k], base_metrics.ScalarMetric)
     }
-    self._writer.write_scalars(step, non_reported_scalars)
+    self._writer.write_scalars(step, non_reported_scalars)  # pyrefly: ignore[bad-argument-type]
     self._writer.flush()
 
     return scalars
@@ -133,7 +133,7 @@ def merge_metrics(
   merged_metrics = {}
   for k in [*a.keys(), *b.keys()]:
     if k in a and k in b:
-      merged_metrics[k] = a[k].merge(b[k])
+      merged_metrics[k] = a[k].merge(b[k])  # pyrefly: ignore[bad-argument-type]
     elif k in a:
       merged_metrics[k] = a[k]
     elif k in b:
@@ -156,4 +156,4 @@ def _localize_and_log_scalars(
 ) -> None:
   """Localizes the metrics from device to host and logs scalars."""
   scalar_metrics = jax.tree.map(_localize, scalar_metrics)
-  summary_writer.write_scalars(step, compute_metrics(scalar_metrics))
+  summary_writer.write_scalars(step, compute_metrics(scalar_metrics))  # pyrefly: ignore[bad-argument-type]

recml/core/ops/hstu_ops.py

@@ -139,7 +139,7 @@ def _apply_mask(
     # need to keep into account the current shard along Q sequence.
 
     if k_in_lanes:
-      assert q_sequence_ref.shape == (bq, NUM_LANES)
+      assert q_sequence_ref.shape == (bq, NUM_LANES)  # pyrefly: ignore[missing-attribute]
 
       k_sequence = k_offset + jax.lax.broadcasted_iota(
           jnp.int32, (bq, k_slice.size), 1
@@ -148,15 +148,15 @@ def _apply_mask(
       repeats, rem = divmod(k_slice.size, NUM_LANES)
       assert rem == 0
       q_sequence = jnp.tile(
-          q_sequence_ref[...], (1, repeats)
+          q_sequence_ref[...], (1, repeats)  # pyrefly: ignore[unsupported-operation]
       )  # [bq, k_slice.size]
     else:
-      assert q_sequence_ref.shape == (NUM_SUBLANES, bq)
+      assert q_sequence_ref.shape == (NUM_SUBLANES, bq)  # pyrefly: ignore[missing-attribute]
 
       k_sequence = k_offset + jax.lax.broadcasted_iota(
           jnp.int32, (k_slice.size, bq), 0
       )
-      q_sequence = q_sequence_ref[:1, :]  # [1, bq]
+      q_sequence = q_sequence_ref[:1, :]  # [1, bq]  # pyrefly: ignore[unsupported-operation]
       q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))
 
     assert q_sequence.shape == k_sequence.shape
@@ -244,7 +244,7 @@ def body(kv_compute_index, _):
         q_sequence_ref=q_sequence_ref,
         q_segment_ids_ref=q_segment_ids_ref,
         kv_segment_ids_ref=kv_segment_ids_ref,
-        k_slice=slice_k,
+        k_slice=slice_k,  # pyrefly: ignore[bad-argument-type]
         # When the iteration space is shrunk (for local attention for example),
         # the kv_index program_id does not correspond to the actual coordinates
         # of the KV data. Make sure to use the 'unshrunk' index (coming from the
@@ -479,7 +479,7 @@ def body(i, _):
         q_sequence_ref,
         q_segment_ids_ref,
         kv_segment_ids_ref,
-        k_slice=slice_k,
+        k_slice=slice_k,  # pyrefly: ignore[bad-argument-type]
         k_offset=j * bkv + i * bkv_compute,
         bq=bq,
         k_in_lanes=False,

recml/core/training/core.py

@@ -161,8 +161,8 @@ def get_iterators(
   """Creates and unpacks the datasets returned by the task."""
   if isinstance(datasets, (iterator.Iterator, tf.data.Dataset)):
     if isinstance(datasets, tf.data.Dataset):
-      datasets = iterator.TFDatasetIterator(datasets)
-    return datasets, {}
+      datasets = iterator.TFDatasetIterator(datasets)  # pyrefly: ignore[bad-assignment]
+    return datasets, {}  # pyrefly: ignore[bad-return]
   elif not isinstance(datasets, tuple) and len(datasets) != 2:
     raise ValueError(
         "Expected `datasets` to be a single dataset or a tuple of training"
@@ -195,7 +195,7 @@ def get_iterators(
 
   if all(isinstance(v, tf.data.Dataset) for v in eval_datasets.values()):
     eval_datasets = {
-        k: iterator.TFDatasetIterator(v) for k, v in eval_datasets.items()
+        k: iterator.TFDatasetIterator(v) for k, v in eval_datasets.items()  # pyrefly: ignore[bad-argument-type]
     }
 
   if not all(isinstance(v, iterator.Iterator) for v in eval_datasets.values()):

recml/core/training/jax_trainer.py

@@ -265,7 +265,7 @@ def update(self, *, grads: PyTree, **kwargs) -> Self:
     tvars_ = dict(zip(self.tvars_paths, self.tvars))
     updates, new_opt_state = self.tx.update(grads_, self.opt_state, tvars_)
     new_tvars_ = optax.apply_updates(tvars_, updates)
-    new_tvars = [new_tvars_[path] for path in self.tvars_paths]
+    new_tvars = [new_tvars_[path] for path in self.tvars_paths]  # pyrefly: ignore[bad-index]
     return self.replace(
         step=self.step + 1,
         tvars=new_tvars,
@@ -881,7 +881,7 @@ def _name(prefix: str, key: str) -> str:
   def _add_optimizer_metrics(opt_state: optax.OptState, prefix: str):
     if isinstance(opt_state, optax.MultiTransformState):
       for name, inner_state in opt_state.inner_states.items():
-        _add_optimizer_metrics(inner_state, _name(prefix, name))
+        _add_optimizer_metrics(inner_state, _name(prefix, name))  # pyrefly: ignore[bad-argument-type]
     elif isinstance(
         opt_state,
         (optax.InjectStatefulHyperparamsState, optax.InjectHyperparamsState),
@@ -892,7 +892,7 @@ def _add_optimizer_metrics(opt_state: optax.OptState, prefix: str):
             and np.prod(hparam.shape) == 1
         ):
           metrics[f"optimizer/{_name(prefix, key)}"] = base_metrics.scalar(
-              hparam
+              hparam  # pyrefly: ignore[bad-argument-type]
           )
     elif isinstance(opt_state, (list, tuple)):
       for opt_state in opt_state:

recml/core/training/keras_trainer.py

@@ -184,7 +184,7 @@ def train_callbacks(self) -> list[keras.callbacks.Callback]:
           ),
       ]
 
-      return callbacks
+      return callbacks  # pyrefly: ignore[bad-return]
 
     return [
         keras.callbacks.TensorBoard(
@@ -310,7 +310,7 @@ def evaluate(self, task: KerasTask) -> core.Logs:
           return_dict=True,
       )
       epoch_dt = time.time() - epoch_start_time
-      steps_per_second = self._steps_per_eval / epoch_dt
+      steps_per_second = self._steps_per_eval / epoch_dt  # pyrefly: ignore[unsupported-operation]
       val_logs = {"val_" + k: v for k, v in history.items()}
       val_logs["val_steps_per_second"] = steps_per_second
       tb_cbk.on_epoch_end(0, val_logs)
@@ -439,7 +439,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
             f" {psutil.Process().memory_info().rss / 1024 ** 2:.1f} MB"
         )
         epoch_dt = time.time() - epoch_start_time
-        steps_per_second = self._steps_per_eval / epoch_dt
+        steps_per_second = self._steps_per_eval / epoch_dt  # pyrefly: ignore[unsupported-operation]
 
         val_logs = {"val_" + k: v for k, v in history.items()}
         val_logs["val_steps_per_second"] = steps_per_second

recml/core/training/partitioning.py

@@ -248,7 +248,7 @@ def partition_init(
       specs = nn.get_partition_spec(abstract_state)
 
       if self.rules is not None:
-        specs = nn.logical_to_mesh(specs, self.rules)
+        specs = nn.logical_to_mesh(specs, self.rules)  # pyrefly: ignore[bad-argument-type]
 
       state_sharding = jax.tree.map(
           lambda x: jax.sharding.NamedSharding(self.mesh, x), specs

recml/core/utils/keras_utils.py

@@ -332,7 +332,7 @@ def restore_keras_checkpoint(
   # TODO(aahil): Look into converging the logic here with the checkpointing
   # logic in KerasOrbaxCheckpointManagerV2.
   checkpointer = ocp.Checkpointer(
-      ocp.CompositeCheckpointHandler(**{
+      ocp.CompositeCheckpointHandler(**{  # pyrefly: ignore[bad-argument-type]
           STATE_CHECKPOINT_KEY: ocp.handlers.PyTreeCheckpointHandler(
               restore_concurrent_gb=96,
           ),
@@ -360,7 +360,7 @@ def restore_keras_checkpoint(
     var._value = restored_var  # pylint: disable=protected-access
 
   if restore_model_epoch:
-    model._initial_epoch = epoch + 1  # pylint: disable=protected-access
+    model._initial_epoch = epoch + 1  # pylint: disable=protected-access  # pyrefly: ignore[unsupported-operation]
   if restore_optimizer_vars and not restore_iterations:
     model.optimizer.iterations.assign(0)
 
@@ -380,7 +380,7 @@ def load_keras_model_config(
 
   json_checkpointer = ocp.Checkpointer(
       ocp.CompositeCheckpointHandler(
-          **{CONFIG_CHECKPOINT_KEY: ocp.handlers.JsonCheckpointHandler()}
+          **{CONFIG_CHECKPOINT_KEY: ocp.handlers.JsonCheckpointHandler()}  # pyrefly: ignore[bad-argument-type]
       )
   )
   cfg = json_checkpointer.restore(
@@ -633,7 +633,7 @@ def restore_keras_model(
     )
 
   checkpointer = ocp.Checkpointer(
-      ocp.CompositeCheckpointHandler(**{
+      ocp.CompositeCheckpointHandler(**{  # pyrefly: ignore[bad-argument-type]
           ORBAX_CHECKPOINT_DEFAULT_KEY: ocp.handlers.PyTreeCheckpointHandler()
       })
   )

recml/examples/DLRM_HSTU/contextual_interleave_preprocessor.py

@@ -140,7 +140,7 @@ def __call__(
     # Prepend contextual embeddings
     if self._max_contextual_seq_len > 0:
       output_seq_embeddings = jnp.concatenate(
-          [contextual_embeddings, output_seq_embeddings], axis=1
+          [contextual_embeddings, output_seq_embeddings], axis=1  # pyrefly: ignore[bad-argument-type]
       )
       contextual_mask = jnp.ones(
           (batch_size, self._max_contextual_seq_len), dtype=jnp.bool_