From 532f9d9af41d19b5c19d1d8be58623b331a092d6 Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Tue, 6 Oct 2026 01:56:00 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 994240996 --- recml/core/metrics/base_metrics.py | 6 +++--- recml/core/metrics/confusion_metrics.py | 6 +++--- recml/core/metrics/reduction_metrics.py | 4 ++-- recml/core/ops/binary_cross_entropy_ops.py | 6 +++--- recml/core/ops/binary_focal_cross_entropy_ops.py | 6 +++--- recml/core/ops/hstu_ops.py | 6 +++--- recml/core/training/core.py | 2 +- recml/examples/dlrm_experiment.py | 2 +- 8 files changed, 19 insertions(+), 19 deletions(-) diff --git a/recml/core/metrics/base_metrics.py b/recml/core/metrics/base_metrics.py index a524afa..ffbc513 100644 --- a/recml/core/metrics/base_metrics.py +++ b/recml/core/metrics/base_metrics.py @@ -31,7 +31,7 @@ # TODO(b/387463777): Consider removing the dependency on CLU metrics longer term # since it's just an interface. -@dataclass_transform(field_specifiers=(struct.field,)) # pytype: disable=not-supported-yet +@dataclass_transform(field_specifiers=(struct.field,)) class Metric(abc.ABC, clu_metrics.Metric, struct.PyTreeNode): """PyTree node representing the state of a metric. @@ -63,7 +63,7 @@ def from_model_output(cls, *args, **kwargs) -> Self: raise NotImplementedError() @abc.abstractmethod - def merge(self, other: Self) -> Self: # pytype: disable=signature-mismatch + def merge(self, other: Self) -> Self: """Merges two metrics. Args: @@ -79,7 +79,7 @@ def merge(self, other: Self) -> Self: # pytype: disable=signature-mismatch """ @abc.abstractmethod - def compute(self) -> Scalar: # pytype: disable=signature-mismatch + def compute(self) -> Scalar: # pyrefly: ignore[bad-override] """Computes the value of the metric. NOTE: This method is almost always called on the host which means that it diff --git a/recml/core/metrics/confusion_metrics.py b/recml/core/metrics/confusion_metrics.py index 7b42e62..75ab808 100644 --- a/recml/core/metrics/confusion_metrics.py +++ b/recml/core/metrics/confusion_metrics.py @@ -67,8 +67,8 @@ def from_model_output( false_negatives=fn, ) - def merge(self, other: Self) -> Self: # pytype: disable=signature-mismatch - return type(self)( # pytype: disable=not-instantiable + def merge(self, other: Self) -> Self: + return type(self)( true_positives=self.true_positives + other.true_positives, true_negatives=self.true_negatives + other.true_negatives, false_positives=self.false_positives + other.false_positives, @@ -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) # pyrefly: ignore[bad-return] + return -np.trapezoid(tp_rate, fp_rate) class PrecisionAtRecall(ConfusionMetric): diff --git a/recml/core/metrics/reduction_metrics.py b/recml/core/metrics/reduction_metrics.py index 509924d..b450b33 100644 --- a/recml/core/metrics/reduction_metrics.py +++ b/recml/core/metrics/reduction_metrics.py @@ -75,7 +75,7 @@ def from_model_output( return cls(total=total) - def merge(self, other: Self) -> Self: # pytype: disable=signature-mismatch + def merge(self, other: Self) -> Self: return type(self)(total=self.total + other.total) def compute(self) -> base_metrics.Scalar: @@ -111,7 +111,7 @@ def from_model_output( return cls(total=total, count=count) - def merge(self, other: Self) -> Self: # pytype: disable=signature-mismatch + def merge(self, other: Self) -> Self: return type(self)( total=self.total + other.total, count=self.count + other.count, diff --git a/recml/core/ops/binary_cross_entropy_ops.py b/recml/core/ops/binary_cross_entropy_ops.py index d92c888..070c63f 100644 --- a/recml/core/ops/binary_cross_entropy_ops.py +++ b/recml/core/ops/binary_cross_entropy_ops.py @@ -1610,7 +1610,7 @@ def _bce_bwd_sharded( vocab_axis_name = emb_spec[0] if emb_spec else None dp_axes = [] - for axis in local_act_spec: # pyrefly: ignore[not-iterable] + for axis in local_act_spec: if axis is not None and axis != vocab_axis_name: dp_axes.append(axis) @@ -1646,13 +1646,13 @@ def _bwd_local_with_reduction(d_loss_, act_, emb_, tgt_): if targets.ndim == replicated_activations.ndim: local_tgt_spec = local_act_spec else: - local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:]) # pyrefly: ignore[unsupported-operation] + local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:]) d_act_replicated, d_emb = jax.shard_map( _bwd_local_with_reduction, mesh=mesh, in_specs=( - jax.sharding.PartitionSpec(*local_act_spec[:-1]), # pyrefly: ignore[unsupported-operation] + jax.sharding.PartitionSpec(*local_act_spec[:-1]), local_act_spec, emb_spec, local_tgt_spec, diff --git a/recml/core/ops/binary_focal_cross_entropy_ops.py b/recml/core/ops/binary_focal_cross_entropy_ops.py index f9d32f0..c0c211a 100644 --- a/recml/core/ops/binary_focal_cross_entropy_ops.py +++ b/recml/core/ops/binary_focal_cross_entropy_ops.py @@ -1187,7 +1187,7 @@ def _focal_bce_bwd_sharded( vocab_axis_name = emb_spec[0] if emb_spec else None dp_axes = [] - for axis in local_act_spec: # pyrefly: ignore[not-iterable] + for axis in local_act_spec: if axis is not None and axis != vocab_axis_name: dp_axes.append(axis) @@ -1223,13 +1223,13 @@ def _bwd_local_with_reduction(d_loss_, act_, emb_, tgt_): if targets.ndim == replicated_activations.ndim: local_tgt_spec = local_act_spec else: - local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:]) # pyrefly: ignore[unsupported-operation] + local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:]) d_act_replicated, d_emb = jax.shard_map( _bwd_local_with_reduction, mesh=mesh, in_specs=( - jax.sharding.PartitionSpec(*local_act_spec[:-1]), # pyrefly: ignore[unsupported-operation] + jax.sharding.PartitionSpec(*local_act_spec[:-1]), local_act_spec, emb_spec, local_tgt_spec, diff --git a/recml/core/ops/hstu_ops.py b/recml/core/ops/hstu_ops.py index 8e76586..bc8c901 100644 --- a/recml/core/ops/hstu_ops.py +++ b/recml/core/ops/hstu_ops.py @@ -160,7 +160,7 @@ def _apply_mask( q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) assert q_sequence.shape == k_sequence.shape - computed_mask = mask_function(q_sequence, k_sequence) # pytype: disable=wrong-arg-count + computed_mask = mask_function(q_sequence, k_sequence) if computed_mask.dtype != jnp.dtype(jnp.bool_): raise ValueError( "Mask function must return a boolean-valued array, but got:" @@ -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, # pyrefly: ignore[bad-argument-type] + k_slice=slice_k, # 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, # pyrefly: ignore[bad-argument-type] + k_slice=slice_k, k_offset=j * bkv + i * bkv_compute, bq=bq, k_in_lanes=False, diff --git a/recml/core/training/core.py b/recml/core/training/core.py index 237993a..75a4d2f 100644 --- a/recml/core/training/core.py +++ b/recml/core/training/core.py @@ -205,7 +205,7 @@ def get_iterators( f" but got {eval_datasets}. You cannot mix both." ) - return train_dataset, eval_datasets # pytype: disable=bad-return-type + return train_dataset, eval_datasets # pyrefly: ignore[bad-return] def get_shape( diff --git a/recml/examples/dlrm_experiment.py b/recml/examples/dlrm_experiment.py index 3c8f5b3..d712335 100644 --- a/recml/examples/dlrm_experiment.py +++ b/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)) # pyrefly: ignore[bad-argument-type] + dataset = tf.data.Dataset.from_tensors((data, label)) dataset = dataset.take(1).repeat() dataset = dataset.prefetch(buffer_size=2048) options = tf.data.Options()