diff --git a/recml/core/metrics/confusion_metrics.py b/recml/core/metrics/confusion_metrics.py index 2e327b1..547c24d 100644 --- a/recml/core/metrics/confusion_metrics.py +++ b/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) diff --git a/recml/core/metrics/reduction_metrics.py b/recml/core/metrics/reduction_metrics.py index c5c1e5f..509924d 100644 --- a/recml/core/metrics/reduction_metrics.py +++ b/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 diff --git a/recml/core/metrics/tools.py b/recml/core/metrics/tools.py index c5e58c9..b6c06f2 100644 --- a/recml/core/metrics/tools.py +++ b/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] diff --git a/recml/core/ops/hstu_ops.py b/recml/core/ops/hstu_ops.py index ab714f7..8e76586 100644 --- a/recml/core/ops/hstu_ops.py +++ b/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, diff --git a/recml/core/training/core.py b/recml/core/training/core.py index 604cbdd..237993a 100644 --- a/recml/core/training/core.py +++ b/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()): diff --git a/recml/core/training/jax_trainer.py b/recml/core/training/jax_trainer.py index 4006888..0a73141 100644 --- a/recml/core/training/jax_trainer.py +++ b/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: diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py index 6eab2a3..c8e121e 100644 --- a/recml/core/training/keras_trainer.py +++ b/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 diff --git a/recml/core/training/partitioning.py b/recml/core/training/partitioning.py index 1d7db15..98d7330 100644 --- a/recml/core/training/partitioning.py +++ b/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 diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index 7355eb7..1109a8f 100644 --- a/recml/core/utils/keras_utils.py +++ b/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() }) ) diff --git a/recml/examples/DLRM_HSTU/contextual_interleave_preprocessor.py b/recml/examples/DLRM_HSTU/contextual_interleave_preprocessor.py index fad1bd3..057ffbb 100644 --- a/recml/examples/DLRM_HSTU/contextual_interleave_preprocessor.py +++ b/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_ diff --git a/recml/examples/DLRM_HSTU/dlrm_hstu.py b/recml/examples/DLRM_HSTU/dlrm_hstu.py index 1da975c..74e2137 100644 --- a/recml/examples/DLRM_HSTU/dlrm_hstu.py +++ b/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, diff --git a/recml/examples/DLRM_HSTU/stu.py b/recml/examples/DLRM_HSTU/stu.py index 82ec7a0..0c328cf 100644 --- a/recml/examples/DLRM_HSTU/stu.py +++ b/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 diff --git a/recml/examples/dlrm_experiment.py b/recml/examples/dlrm_experiment.py index a9a3b23..3c8f5b3 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)) + 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 diff --git a/recml/layers/linen/sparsecore.py b/recml/layers/linen/sparsecore.py index c98d9f2..74e13ac 100644 --- a/recml/layers/linen/sparsecore.py +++ b/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'].")