Internal change

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_

recml/examples/DLRM_HSTU/dlrm_hstu.py

@@ -230,7 +230,7 @@ def setup(self):
     action_encoder = ActionEncoder(
         action_embedding_dim=hstu_config.hstu_transducer_embedding_dim,
         action_feature_name=hstu_config.uih_weight_feature_name,
-        action_weights=hstu_config.action_weights,
+        action_weights=hstu_config.action_weights,  # pyrefly: ignore[bad-argument-type]
         watchtime_feature_name=hstu_config.watchtime_feature_name,
         watchtime_to_action_thresholds_and_weights=hstu_config.watchtime_to_action_thresholds_and_weights,
     )
@@ -321,7 +321,7 @@ def mlp_fn(
     self._hstu_transducer = HSTUTransducer(
         stu_module=stu_module,
         input_preprocessor=preprocessor,
-        output_postprocessor_cls=postproc_cls,
+        output_postprocessor_cls=postproc_cls,  # pyrefly: ignore[bad-argument-type]
         input_dropout_ratio=hstu_config.hstu_input_dropout_ratio,
         positional_encoder=positional_encoder,
         return_full_embeddings=False,

recml/examples/DLRM_HSTU/stu.py

@@ -262,22 +262,22 @@ def hstu_preprocess_and_attention(
             self.config.max_decode_length,
             self.hidden_dim,
         )
-        self.cached_key.value = jnp.zeros(k_cache_shape, k.dtype)
-        self.cached_value.value = jnp.zeros(v_cache_shape, v.dtype)
+        self.cached_key.value = jnp.zeros(k_cache_shape, k.dtype)  # pyrefly: ignore[bad-argument-type]
+        self.cached_value.value = jnp.zeros(v_cache_shape, v.dtype)  # pyrefly: ignore[bad-argument-type]
 
       if self.is_mutable_collection('cache'):
         k_cache = jax.lax.dynamic_update_slice(
-            self.cached_key.value,
-            k.astype(self.cached_key.value.dtype),
+            self.cached_key.value,  # pyrefly: ignore[bad-argument-type]
+            k.astype(self.cached_key.value.dtype),  # pyrefly: ignore[missing-attribute]
             (0, 0, cache_index, 0),
         )
         v_cache = jax.lax.dynamic_update_slice(
-            self.cached_value.value,
-            v.astype(self.cached_value.value.dtype),
+            self.cached_value.value,  # pyrefly: ignore[bad-argument-type]
+            v.astype(self.cached_value.value.dtype),  # pyrefly: ignore[missing-attribute]
             (0, 0, cache_index, 0),
         )
-        self.cached_key.value = k_cache
-        self.cached_value.value = v_cache
+        self.cached_key.value = k_cache  # pyrefly: ignore[bad-argument-type]
+        self.cached_value.value = v_cache  # pyrefly: ignore[bad-argument-type]
         self.cache_index.value = cache_index + seq_len
         k = k_cache
         v = v_cache

recml/examples/dlrm_experiment.py

@@ -215,7 +215,7 @@ def make(self) -> tf.data.Dataset:
 
     label = np.random.randint(0, 2, size=(batch_size,))
 
-    dataset = tf.data.Dataset.from_tensors((data, label))
+    dataset = tf.data.Dataset.from_tensors((data, label))  # pyrefly: ignore[bad-argument-type]
     dataset = dataset.take(1).repeat()
     dataset = dataset.prefetch(buffer_size=2048)
     options = tf.data.Options()
@@ -234,7 +234,7 @@ class PredictionTask(recml.JaxTask):
   model: DLRMModel
   optimizer: recml.Factory[optax.GradientTransformation]
 
-  def create_datasets(self) -> tuple[recml.data.Iterator, recml.data.Iterator]:
+  def create_datasets(self) -> tuple[recml.data.Iterator, recml.data.Iterator]:  # pyrefly: ignore[bad-override]
     global_batch_size = self.train_data.global_batch_size
     train_iter = recml.data.TFDatasetIterator(
         dataset=self.train_data.make(),
@@ -263,8 +263,8 @@ def train_step(
 
     def _loss_fn(params: jt.PyTree) -> tuple[jt.Scalar, jt.Array]:
       logits = self.model.apply(params, inputs, training=True)
-      loss = jnp.mean(optax.sigmoid_binary_cross_entropy(logits, label), axis=0)
-      return loss, logits
+      loss = jnp.mean(optax.sigmoid_binary_cross_entropy(logits, label), axis=0)  # pyrefly: ignore[bad-argument-type]
+      return loss, logits  # pyrefly: ignore[bad-return]
 
     grad_fn = jax.value_and_grad(_loss_fn, has_aux=True, allow_int=True)
     (loss, logits), grads = grad_fn(state.params)
@@ -285,15 +285,15 @@ def eval_step(
   ) -> Mapping[str, recml.Metric]:
     inputs, label = batch
     logits = self.model.apply(state.params, inputs, training=False)
-    loss = jnp.mean(optax.sigmoid_binary_cross_entropy(logits, label), axis=0)
+    loss = jnp.mean(optax.sigmoid_binary_cross_entropy(logits, label), axis=0)  # pyrefly: ignore[bad-argument-type]
 
     metrics = {
         'loss': recml.metrics.mean(loss),
-        'accuracy': recml.metrics.binary_accuracy(label, logits, threshold=0.0),
-        'auc': recml.metrics.aucpr(label, logits, from_logits=True),
-        'aucroc': recml.metrics.aucroc(label, logits, from_logits=True),
+        'accuracy': recml.metrics.binary_accuracy(label, logits, threshold=0.0),  # pyrefly: ignore[bad-argument-type]
+        'auc': recml.metrics.aucpr(label, logits, from_logits=True),  # pyrefly: ignore[bad-argument-type]
+        'aucroc': recml.metrics.aucroc(label, logits, from_logits=True),  # pyrefly: ignore[bad-argument-type]
         'label/mean': recml.metrics.mean(label),
-        'prediction/mean': recml.metrics.mean(jax.nn.sigmoid(logits)),
+        'prediction/mean': recml.metrics.mean(jax.nn.sigmoid(logits)),  # pyrefly: ignore[bad-argument-type]
     }
     return metrics
 

recml/layers/linen/sparsecore.py

@@ -551,7 +551,7 @@ def cpu_lookup(
       activation = tf.nn.embedding_lookup(tables[name], feature)
 
       if spec.max_sequence_length is None:
-        activation = _reduce(activation, weight, spec.combiner)
+        activation = _reduce(activation, weight, spec.combiner)  # pyrefly: ignore[bad-argument-type]
 
       activations[name] = activation
     else:
@@ -591,7 +591,7 @@ def _reduce(
     weight_sum = tf.reduce_sum(weights, axis=-2)
     out = tf.math.divide_no_nan(out, weight_sum)
   elif combiner == 'sqrtn':
-    weight_sum = tf.math.sqrt(tf.reduce_sum(weights**2, axis=-2))
+    weight_sum = tf.math.sqrt(tf.reduce_sum(weights**2, axis=-2))  # pyrefly: ignore[unsupported-operation]
     out = tf.math.divide_no_nan(out, weight_sum)
   else:
     raise ValueError("`combiner` must be one of ['mean', 'sqrtn', 'sum'].")