From ce73bb7f034fc377545e89d69f5ae94ee255e72f Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Tue, 7 Jul 2026 04:16:03 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 943808612 --- .../checkpoint_utils.py | 8 +-- .../input_utils.py | 2 +- .../preprocess_utils.py | 4 +- .../train_utils.py | 6 +-- .../drug_cardiotoxicity/augmentation_utils.py | 6 +-- .../drug_cardiotoxicity/deterministic.py | 12 ++--- baselines/drug_cardiotoxicity/sngp.py | 14 ++--- baselines/drug_cardiotoxicity/utils.py | 4 +- .../cifar_pi/distill_pi.py | 32 ++++++------ .../privileged_information/cifar_pi/no_pi.py | 26 +++++----- .../cifar_pi/pi_full_marginalization.py | 36 ++++++------- .../privileged_information/cifar_pi/tram.py | 52 +++++++++---------- .../imagenet_pi/distill_pi.py | 20 +++---- .../imagenet_pi/no_pi.py | 12 ++--- .../imagenet_pi/pi_full_marginalization.py | 16 +++--- .../imagenet_pi/tram.py | 16 +++--- .../t5/data/deepbank/graph_linear_utils.py | 8 +-- baselines/t5/data/deepbank/graph_utils.py | 8 +-- baselines/t5/data/deepbank/lispress_utils.py | 2 +- baselines/t5/data/deepbank/penman_utils.py | 2 +- baselines/t5/data/tasks/deepbank.py | 6 +-- baselines/t5/data/tasks/mnli.py | 4 +- baselines/t5/data/tasks/nalue.py | 8 +-- baselines/t5/data/tasks/toxic_comments.py | 4 +- baselines/t5/models/be_models.py | 4 +- baselines/t5/models/gp_models.py | 2 +- baselines/t5/models/heteroscedastic_models.py | 2 +- baselines/t5/models/models.py | 32 ++++++------ experimental/language_structure/psl/data.py | 10 ++-- .../language_structure/psl/psl_model.py | 16 +++--- .../psl/psl_model_dstc_synthetic.py | 2 +- .../psl/psl_model_multiwoz_test_util.py | 6 +-- experimental/multimodal/checkpoint_utils.py | 12 ++--- experimental/multimodal/deterministic.py | 8 +-- experimental/multimodal/input_utils.py | 8 +-- experimental/multimodal/preprocess_utils.py | 2 +- .../models/genomics_cnn.py | 8 +-- .../single_model_uncertainty/models/models.py | 2 +- .../single_model_uncertainty/models/util.py | 4 +- .../models/wide_resnet.py | 16 +++--- 40 files changed, 221 insertions(+), 221 deletions(-) diff --git a/baselines/diabetic_retinopathy_detection/checkpoint_utils.py b/baselines/diabetic_retinopathy_detection/checkpoint_utils.py index 5b2fde1bc..b6b358a2f 100644 --- a/baselines/diabetic_retinopathy_detection/checkpoint_utils.py +++ b/baselines/diabetic_retinopathy_detection/checkpoint_utils.py @@ -205,7 +205,7 @@ def checkpoint_trained_model( accum_train_time=checkpoint_data.accumulated_train_time), ) if checkpoint_data.fixed_model_states is not None: - tree["states"] = checkpoint_data.fixed_model_states + tree["states"] = checkpoint_data.fixed_model_states # pyrefly: ignore[bad-assignment] save_checkpoint(tree, path, step_for_copy) @@ -216,7 +216,7 @@ def _flatten_jax_params_dict(d: Params, parent_key: str = "", for k, v in d.items(): path = parent_key + sep + k if parent_key else k if isinstance(v, collections.abc.Mapping): - items.extend(_flatten_jax_params_dict(v, path, sep=sep).items()) + items.extend(_flatten_jax_params_dict(v, path, sep=sep).items()) # pyrefly: ignore[bad-argument-type] else: items.append((path, v)) @@ -411,7 +411,7 @@ def maybe_load_checkpoint(train_loop_rngs: jnp.ndarray, checkpoint_tree = {"opt": init_optimizer, "extra": checkpoint_extra} if init_fixed_model_states is not None: - checkpoint_tree["states"] = init_fixed_model_states + checkpoint_tree["states"] = init_fixed_model_states # pyrefly: ignore[bad-assignment] checkpoint = load_checkpoint(checkpoint_tree, resume_checkpoint_path) optimizer, checkpoint_extra = checkpoint["opt"], checkpoint["extra"] fixed_model_states = checkpoint.get("states", None) @@ -442,7 +442,7 @@ def maybe_load_checkpoint(train_loop_rngs: jnp.ndarray, return CheckpointData( optimizer=optimizer, fixed_model_states=fixed_model_states, - train_loop_rngs=checkpoint_extra["rngs_loop"], + train_loop_rngs=checkpoint_extra["rngs_loop"], # pyrefly: ignore[bad-argument-type] accumulated_train_time=checkpoint_extra["accum_train_time"]) diff --git a/baselines/diabetic_retinopathy_detection/input_utils.py b/baselines/diabetic_retinopathy_detection/input_utils.py index 8d28982b0..d8d3576f0 100644 --- a/baselines/diabetic_retinopathy_detection/input_utils.py +++ b/baselines/diabetic_retinopathy_detection/input_utils.py @@ -206,7 +206,7 @@ def get_data( dataset_builder = _get_dataset_builder(dataset, data_dir) if rng is not None: - rng = jax.random.fold_in(rng, process_index) # Derive RNG for this process. + rng = jax.random.fold_in(rng, process_index) # Derive RNG for this process. # pyrefly: ignore[bad-argument-type] process_split = _get_process_split( split, diff --git a/baselines/diabetic_retinopathy_detection/preprocess_utils.py b/baselines/diabetic_retinopathy_detection/preprocess_utils.py index 598bbe942..e22ad2874 100644 --- a/baselines/diabetic_retinopathy_detection/preprocess_utils.py +++ b/baselines/diabetic_retinopathy_detection/preprocess_utils.py @@ -81,7 +81,7 @@ def __call__(self, features: Features) -> Features: # Perform binarization using given threshold labels = features["label"] - labels = tf.cast(labels > highest_negative_class, tf.int32) + labels = tf.cast(labels > highest_negative_class, tf.int32) # pyrefly: ignore[unsupported-operation] features[self.key_result or self.key] = decoded_image features["labels"] = labels del features["label"] @@ -413,7 +413,7 @@ def __call__(self, features: Features) -> Features: # than using tf.one_hot followed by tf.reduce_max; we tested. labels = features[self.key] if labels.shape.rank > 0 and self.multi: # pytype: disable=attribute-error - x = tf.scatter_nd(labels[:, None], tf.ones(tf.shape(labels)[0]), + x = tf.scatter_nd(labels[:, None], tf.ones(tf.shape(labels)[0]), # pyrefly: ignore[bad-index] (self.depth,)) x = tf.clip_by_value(x, 0, 1) * (self.on - self.off) + self.off else: diff --git a/baselines/diabetic_retinopathy_detection/train_utils.py b/baselines/diabetic_retinopathy_detection/train_utils.py index 6ee946fe3..da31f86a9 100644 --- a/baselines/diabetic_retinopathy_detection/train_utils.py +++ b/baselines/diabetic_retinopathy_detection/train_utils.py @@ -260,7 +260,7 @@ def hms(s): return f"{h:.0f}h{m:.0f}m" # Seconds intentionally omitted. # Progress note with "global" full-program average timings - dt = now - self.start_time # Time since process start. + dt = now - self.start_time # Time since process start. # pyrefly: ignore[unsupported-operation] steps_done = step - self.first_step steps_todo = self.total_steps - step self.note = f"Steps:{step}/{self.total_steps} [{step/self.total_steps:.1%}]" @@ -270,7 +270,7 @@ def hms(s): timing_measurements = {} # Measurement with micro-timings of current training steps speed. - dt = now - self.prev_time - self.paused_time # Time between ticks. + dt = now - self.prev_time - self.paused_time # Time between ticks. # pyrefly: ignore[unsupported-operation] ds = step - self.prev_step # Steps between ticks. ncores = jax.device_count() # Global device count. timing_measurements["img/sec/core"] = self.global_bs * ds / dt / ncores @@ -294,5 +294,5 @@ def pause(self): def resume(self): """Resumes the time measurement.""" - self.paused_time += time.time() - self.pause_start + self.paused_time += time.time() - self.pause_start # pyrefly: ignore[bad-assignment, unsupported-operation] self.pause_start = None diff --git a/baselines/drug_cardiotoxicity/augmentation_utils.py b/baselines/drug_cardiotoxicity/augmentation_utils.py index 85ce5369d..68f8f415d 100644 --- a/baselines/drug_cardiotoxicity/augmentation_utils.py +++ b/baselines/drug_cardiotoxicity/augmentation_utils.py @@ -330,7 +330,7 @@ def _perturb_edges_helper( perturbation_values_for_pair_mask = drop_values_for_pair_mask else: edges_to_perturb = tf.concat( - (idx_edges_to_drop, idx_bidirectional_edges_to_add), axis=0) + (idx_edges_to_drop, idx_bidirectional_edges_to_add), axis=0) # pyrefly: ignore[unbound-name] add_values_for_pair_mask = tf.ones( tf.shape(idx_bidirectional_edges_to_add)[0]) perturbation_values_for_pair_mask = tf.concat( @@ -344,11 +344,11 @@ def _perturb_edges_helper( else: if self.initialize_edge_features_randomly: add_values_for_pairs = tf.random.uniform( - (tf.shape(idx_bidirectional_edges_to_add)[0], + (tf.shape(idx_bidirectional_edges_to_add)[0], # pyrefly: ignore[unbound-name] tf.shape(pairs)[-1])) else: add_values_for_pairs = tf.concat( - (features_of_edges_to_drop, features_of_edges_to_drop), axis=0) + (features_of_edges_to_drop, features_of_edges_to_drop), axis=0) # pyrefly: ignore[unbound-name] perturbation_values_for_pairs = tf.concat( (drop_values_for_pairs, add_values_for_pairs), axis=0) diff --git a/baselines/drug_cardiotoxicity/deterministic.py b/baselines/drug_cardiotoxicity/deterministic.py index 04e9dcf74..4197e35b8 100644 --- a/baselines/drug_cardiotoxicity/deterministic.py +++ b/baselines/drug_cardiotoxicity/deterministic.py @@ -249,7 +249,7 @@ def step_fn(inputs): negative_log_likelihood = tf.reduce_mean( tf.keras.losses.categorical_crossentropy(labels, probs)) - metrics[f'{dataset_name}/negative_log_likelihood'].update_state( + metrics[f'{dataset_name}/negative_log_likelihood'].update_state( # pyrefly: ignore[missing-argument] negative_log_likelihood) metrics[f'{dataset_name}/accuracy'].update_state(labels, probs) metrics[f'{dataset_name}/roc_auc'].update_state( @@ -275,8 +275,8 @@ def step_fn(inputs): logging.info('Starting to run epoch: %s', epoch) train_step(train_iterator) - current_step = (epoch + 1) * params.steps_per_epoch - max_steps = params.steps_per_epoch * params.num_epochs + current_step = (epoch + 1) * params.steps_per_epoch # pyrefly: ignore[unsupported-operation] + max_steps = params.steps_per_epoch * params.num_epochs # pyrefly: ignore[unsupported-operation] time_elapsed = time.time() - start_time steps_per_sec = float(current_step) / time_elapsed eta_seconds = (max_steps - current_step) / steps_per_sec @@ -325,11 +325,11 @@ def main(argv: Sequence[str]): strategy = tf.distribute.MirroredStrategy() train_dataset, steps_per_epoch = utils.load_dataset(FLAGS.data_dir, - tfds.Split.TRAIN, + tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] FLAGS.batch_size) eval_identifiers = ['tune', 'test1', 'test2'] - splits = [tfds.Split.VALIDATION, tfds.Split.TEST, tfds.Split('test2')] + splits = [tfds.Split.VALIDATION, tfds.Split.TEST, tfds.Split('test2')] # pyrefly: ignore[missing-attribute] eval_datasets, steps_per_eval = utils.load_eval_datasets( eval_identifiers, splits, FLAGS.data_dir, FLAGS.batch_size) @@ -380,7 +380,7 @@ def main(argv: Sequence[str]): strategy=strategy, summary_writer=summary_writer, loss_type=FLAGS.loss_type, - graph_augmenter=graph_augmenter) + graph_augmenter=graph_augmenter) # pyrefly: ignore[bad-argument-type] if __name__ == '__main__': diff --git a/baselines/drug_cardiotoxicity/sngp.py b/baselines/drug_cardiotoxicity/sngp.py index 2a889b109..edb370e88 100644 --- a/baselines/drug_cardiotoxicity/sngp.py +++ b/baselines/drug_cardiotoxicity/sngp.py @@ -181,7 +181,7 @@ def per_replica_train_step_fn(inputs): sample_weights = 1 with tf.GradientTape() as tape: - probs = model(features, training=True) + probs = model(features, training=True) # pyrefly: ignore[not-callable] negative_log_likelihood = tf.reduce_mean( tf.keras.losses.categorical_crossentropy(labels, probs) * sample_weights) @@ -215,11 +215,11 @@ def per_replica_eval_step_fn(inputs, dataset_name): else: features, labels = inputs - probs = model(features, training=False) + probs = model(features, training=False) # pyrefly: ignore[not-callable] negative_log_likelihood = tf.reduce_mean( tf.keras.losses.categorical_crossentropy(labels, probs)) - metrics[f'{dataset_name}/negative_log_likelihood'].update_state( + metrics[f'{dataset_name}/negative_log_likelihood'].update_state( # pyrefly: ignore[missing-argument] negative_log_likelihood) metrics[f'{dataset_name}/accuracy'].update_state(labels, probs) metrics[f'{dataset_name}/roc_auc'].update_state(labels[:, 1], probs[:, 1]) @@ -254,8 +254,8 @@ def distributed_eval_step(iterator, dataset_name, num_steps): logging.info('Starting to run epoch: %s', epoch) distributed_train_step(train_iterator) - current_step = (epoch + 1) * params.steps_per_epoch - max_steps = params.steps_per_epoch * params.num_epochs + current_step = (epoch + 1) * params.steps_per_epoch # pyrefly: ignore[unsupported-operation] + max_steps = params.steps_per_epoch * params.num_epochs # pyrefly: ignore[unsupported-operation] time_elapsed = time.time() - start_time steps_per_sec = float(current_step) / time_elapsed eta_seconds = (max_steps - current_step) / steps_per_sec @@ -305,11 +305,11 @@ def main(argv: Sequence[str]): strategy = tf.distribute.MirroredStrategy() train_dataset, steps_per_epoch = utils.load_dataset(FLAGS.data_dir, - tfds.Split.TRAIN, + tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] FLAGS.batch_size) eval_identifiers = ['tune', 'test1', 'test2'] - splits = [tfds.Split.VALIDATION, tfds.Split.TEST, tfds.Split('test2')] + splits = [tfds.Split.VALIDATION, tfds.Split.TEST, tfds.Split('test2')] # pyrefly: ignore[missing-attribute] eval_datasets, steps_per_eval = utils.load_eval_datasets( eval_identifiers, splits, FLAGS.data_dir, FLAGS.batch_size) diff --git a/baselines/drug_cardiotoxicity/utils.py b/baselines/drug_cardiotoxicity/utils.py index 90e5612e2..ce40a40b3 100644 --- a/baselines/drug_cardiotoxicity/utils.py +++ b/baselines/drug_cardiotoxicity/utils.py @@ -112,11 +112,11 @@ def get_metric_result_value(metric): def load_dataset(data_dir, split, batch_size): """Loads a single dataset with specific split.""" known_splits = [ - tfds.Split.TRAIN, tfds.Split.VALIDATION, tfds.Split.TEST, + tfds.Split.TRAIN, tfds.Split.VALIDATION, tfds.Split.TEST, # pyrefly: ignore[missing-attribute] tfds.Split('test2') ] if split in known_splits: - is_training = split == tfds.Split.TRAIN + is_training = split == tfds.Split.TRAIN # pyrefly: ignore[missing-attribute] else: raise ValueError( 'Received ambiguous split {}, must set is_training for splits other ' diff --git a/baselines/privileged_information/cifar_pi/distill_pi.py b/baselines/privileged_information/cifar_pi/distill_pi.py index f5db4e4bc..266e9a183 100644 --- a/baselines/privileged_information/cifar_pi/distill_pi.py +++ b/baselines/privileged_information/cifar_pi/distill_pi.py @@ -288,8 +288,8 @@ def main(argv): steps_per_validation = validation_builder.num_examples // batch_size clean_test_builder = ub.datasets.get( 'cifar10' if FLAGS.dataset in ['cifar10n', 'cifar10h'] else 'cifar100', - split=tfds.Split.TEST - if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, + split=tfds.Split.TEST # pyrefly: ignore[missing-attribute] + if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] data_dir=data_dir, drop_remainder=FLAGS.drop_remainder_for_eval, is_training=False) @@ -328,7 +328,7 @@ def main(argv): f'{FLAGS.dataset}_corrupted', corruption_type=corruption_type, severity=severity, - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] data_dir=data_dir).load(batch_size=batch_size) test_datasets[f'{corruption_type}_{severity}'] = ( strategy.experimental_distribute_dataset(dataset)) @@ -375,7 +375,7 @@ def annotator_labels_encoding(example): 'random_pi': lambda e: e['pi_features']['random_pi'], } privileged_information_fn = pi_utils.get_privileged_information_fn( - pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) + pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) # pyrefly: ignore[bad-argument-type] pi_shape = privileged_information_fn(dummy_example).shape[1:] @@ -470,12 +470,12 @@ def create_optimizer(epochs): rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins), }) if FLAGS.eval_on_ood: - ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names) + ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names) # pyrefly: ignore[unbound-name] metrics.update(ood_metrics) if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, 6): - for corruption in corruption_types: + for corruption in corruption_types: # pyrefly: ignore[unbound-name] dataset_name = '{0}_{1}'.format(corruption, intensity) corrupt_metrics[f'distill/test/nll_{dataset_name}'] = ( tf.keras.metrics.Mean()) @@ -590,7 +590,7 @@ def update_train_metrics(labels, sample_weight = None if noise_split == '_clean': - sample_weight = 1. - noisy_idx + sample_weight = 1. - noisy_idx # pyrefly: ignore[unsupported-operation] elif noise_split == '_noisy': sample_weight = noisy_idx metrics[f'{training_phase}/train/ece{noise_split}'].add_batch( @@ -613,7 +613,7 @@ def step_fn(inputs): noisy_idx = pi_utils.find_noisy_annotators(inputs) with tf.GradientTape() as tape: - logits = teacher_model((images, privileged_information), training=True) + logits = teacher_model((images, privileged_information), training=True) # pyrefly: ignore[not-callable] # Flatten the annotator axis. logits = pi_utils.flatten_annotator_axis(logits) @@ -668,9 +668,9 @@ def step_fn(inputs): noisy_idx = pi_utils.find_noisy_annotators(inputs) with tf.GradientTape() as tape: - teacher_logits = teacher_model((images, privileged_information), + teacher_logits = teacher_model((images, privileged_information), # pyrefly: ignore[not-callable] training=False) - logits = distill_model(images, training=True) + logits = distill_model(images, training=True) # pyrefly: ignore[not-callable] # teacher_logits: (batch_size, num_annotators, num_classes) -> # (batch_size * num_annotators, num_classes) @@ -755,7 +755,7 @@ def step_fn(inputs): else: labels = inputs['clean_labels'] - logits = distill_model(images, training=False) + logits = distill_model(images, training=False) # pyrefly: ignore[not-callable] # pytype: disable=attribute-error logits = pi_utils.repeat_across_annotators( logits, @@ -764,7 +764,7 @@ def step_fn(inputs): # pytype: enable=attribute-error logits = pi_utils.flatten_annotator_axis(logits) if use_annotator_labels: - logits = tf.gather(logits, non_empty_indices) + logits = tf.gather(logits, non_empty_indices) # pyrefly: ignore[unbound-name] probs = tf.nn.softmax(logits) negative_log_likelihood = tf.reduce_mean( @@ -789,7 +789,7 @@ def step_fn(inputs): images = inputs['features'] labels = inputs['labels'] - logits = distill_model(images, training=False) + logits = distill_model(images, training=False) # pyrefly: ignore[not-callable] probs = tf.nn.softmax(logits) negative_log_likelihood = tf.reduce_mean( @@ -941,19 +941,19 @@ def write_metrics_to_summary(total_results, epoch): logging.info('Done with testing on %s', dataset_name) if FLAGS.eval_on_ood: - for ood_dataset_name, ood_dataset in ood_datasets.items(): + for ood_dataset_name, ood_dataset in ood_datasets.items(): # pyrefly: ignore[unbound-name] ood_iterator = iter(ood_dataset) logging.info('Calculating OOD on dataset %s', ood_dataset_name) logging.info('Running OOD eval at epoch: %s', epoch) test_step(ood_iterator, 'test', ood_dataset_name, - steps_per_ood[ood_dataset_name]) + steps_per_ood[ood_dataset_name]) # pyrefly: ignore[unbound-name] logging.info('Done with OOD eval on %s', ood_dataset_name) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): - corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, + corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, # pyrefly: ignore[unbound-name] corruption_types) logging.info('Distillation Loss: %.4f, Accuracy: %.2f%%', diff --git a/baselines/privileged_information/cifar_pi/no_pi.py b/baselines/privileged_information/cifar_pi/no_pi.py index 712fb6d8e..9518c604e 100644 --- a/baselines/privileged_information/cifar_pi/no_pi.py +++ b/baselines/privileged_information/cifar_pi/no_pi.py @@ -296,8 +296,8 @@ def main(argv): steps_per_validation = validation_builder.num_examples // batch_size clean_test_builder = ub.datasets.get( 'cifar10' if FLAGS.dataset in ['cifar10n', 'cifar10h'] else 'cifar100', - split=tfds.Split.TEST - if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, + split=tfds.Split.TEST # pyrefly: ignore[missing-attribute] + if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] data_dir=data_dir, drop_remainder=FLAGS.drop_remainder_for_eval, is_training=False) @@ -335,7 +335,7 @@ def main(argv): f'{FLAGS.dataset}_corrupted', corruption_type=corruption_type, severity=severity, - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] data_dir=data_dir).load(batch_size=batch_size) test_datasets[f'{corruption_type}_{severity}'] = ( strategy.experimental_distribute_dataset(dataset)) @@ -415,12 +415,12 @@ def main(argv): rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins), }) if FLAGS.eval_on_ood: - ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names) + ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names) # pyrefly: ignore[unbound-name] metrics.update(ood_metrics) if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, 6): - for corruption in corruption_types: + for corruption in corruption_types: # pyrefly: ignore[unbound-name] dataset_name = '{0}_{1}'.format(corruption, intensity) corrupt_metrics['test/nll_{}'.format(dataset_name)] = ( tf.keras.metrics.Mean()) @@ -486,7 +486,7 @@ def step_fn(inputs): # We take just 1 augmented image from the returned augmented images. images = images[:, 1, ...] with tf.GradientTape() as tape: - logits = model(images, training=True) + logits = model(images, training=True) # pyrefly: ignore[not-callable] if FLAGS.use_annotator_labels: # pytype: disable=attribute-error logits = pi_utils.repeat_across_annotators( @@ -495,7 +495,7 @@ def step_fn(inputs): .num_annotators_per_example_and_step) # pytype: enable=attribute-error logits = pi_utils.flatten_annotator_axis(logits) - logits = tf.gather(logits, non_empty_indices) + logits = tf.gather(logits, non_empty_indices) # pyrefly: ignore[unbound-name] if FLAGS.label_smoothing == 0.: negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( @@ -532,7 +532,7 @@ def step_fn(inputs): metrics[f'train/ece{noise_split}'].add_batch( probs, label=labels, sample_weight=sample_weight) metrics[f'train/accuracy{noise_split}'].update_state( - labels, logits, sample_weight=sample_weight) + labels, logits, sample_weight=sample_weight) # pyrefly: ignore[bad-keyword-argument] metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( @@ -548,7 +548,7 @@ def test_step(iterator, dataset_split, dataset_name, num_steps): def step_fn(inputs): """Per-Replica StepFn.""" images = inputs['features'] - logits = model(images, training=False) + logits = model(images, training=False) # pyrefly: ignore[not-callable] probs = tf.nn.softmax(logits) if dataset_split == 'validation_noisy': @@ -655,19 +655,19 @@ def step_fn(inputs): logging.info('Done with testing on %s.', dataset_name) if FLAGS.eval_on_ood: - for ood_dataset_name, ood_dataset in ood_datasets.items(): + for ood_dataset_name, ood_dataset in ood_datasets.items(): # pyrefly: ignore[unbound-name] ood_iterator = iter(ood_dataset) logging.info('Calculating OOD on dataset %s.', ood_dataset_name) logging.info('Running OOD eval at epoch: %s.', epoch) test_step(ood_iterator, 'test', ood_dataset_name, - steps_per_ood[ood_dataset_name]) + steps_per_ood[ood_dataset_name]) # pyrefly: ignore[unbound-name] logging.info('Done with OOD eval on %s.', ood_dataset_name) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): - corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, + corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, # pyrefly: ignore[unbound-name] corruption_types) logging.info('Train Loss: %.4f, Accuracy: %.2f%%', @@ -692,7 +692,7 @@ def step_fn(inputs): metric.reset_states() if FLAGS.corruptions_interval > 0: - for metric in corrupt_metrics.values(): + for metric in corrupt_metrics.values(): # pyrefly: ignore[unbound-name] metric.reset_states() if (FLAGS.checkpoint_interval > 0 and diff --git a/baselines/privileged_information/cifar_pi/pi_full_marginalization.py b/baselines/privileged_information/cifar_pi/pi_full_marginalization.py index 7d302fe28..54216b334 100644 --- a/baselines/privileged_information/cifar_pi/pi_full_marginalization.py +++ b/baselines/privileged_information/cifar_pi/pi_full_marginalization.py @@ -338,8 +338,8 @@ def main(argv): steps_per_validation = validation_builder.num_examples // batch_size clean_test_builder = ub.datasets.get( 'cifar10' if FLAGS.dataset in ['cifar10n', 'cifar10h'] else 'cifar100', - split=tfds.Split.TEST - if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, + split=tfds.Split.TEST # pyrefly: ignore[missing-attribute] + if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] data_dir=data_dir, drop_remainder=FLAGS.drop_remainder_for_eval, is_training=False) @@ -380,7 +380,7 @@ def main(argv): f'{FLAGS.dataset}_corrupted', corruption_type=corruption_type, severity=severity, - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] data_dir=data_dir).load(batch_size=batch_size) test_datasets[f'{corruption_type}_{severity}'] = ( strategy.experimental_distribute_dataset(dataset)) @@ -427,7 +427,7 @@ def annotator_labels_encoding(example): 'random_pi': lambda e: e['pi_features']['random_pi'], } privileged_information_fn = pi_utils.get_privileged_information_fn( - pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) + pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) # pyrefly: ignore[bad-argument-type] pi_shape = privileged_information_fn(dummy_example).shape[1:] @@ -586,7 +586,7 @@ def step_fn(inputs): # We take just 1 augmented image from the returned augmented images. images = images[:, 1, ...] with tf.GradientTape() as tape: - (logits_pi, probs_fm) = model( + (logits_pi, probs_fm) = model( # pyrefly: ignore[not-callable] (images, privileged_information, batched_pi_mc_samples), training=True) @@ -644,7 +644,7 @@ def step_fn(inputs): } for setting in ['pi', 'fm']: - metrics[f'train/negative_log_likelihood_{setting}'].update_state( + metrics[f'train/negative_log_likelihood_{setting}'].update_state( # pyrefly: ignore[missing-attribute] stats[setting]['negative_log_likelihood']) for noise_split in ['', '_clean', '_noisy']: sample_weight = None @@ -656,7 +656,7 @@ def step_fn(inputs): stats[setting]['probs'], label=labels, sample_weight=sample_weight) - metrics[f'train/accuracy_{setting}{noise_split}'].update_state( + metrics[f'train/accuracy_{setting}{noise_split}'].update_state( # pyrefly: ignore[missing-attribute] labels, stats[setting]['outputs'], sample_weight=sample_weight) # pylint: disable=redundant-keyword-arg metrics['train/loss'].update_state(loss) @@ -705,7 +705,7 @@ def step_fn(inputs): tf.expand_dims(pi_mc_samples, axis=0), multiples=[batch_size, 1, 1, 1] ) - _, probs_fm = model( + _, probs_fm = model( # pyrefly: ignore[not-callable] (images, random_pi, batched_pi_mc_samples), training=False ) probs_fm = pi_utils.flatten_annotator_axis(probs_fm) @@ -719,10 +719,10 @@ def step_fn(inputs): validation_type = ( 'validation_noisy' if use_annotator_labels else 'validation_clean' ) - metrics[f'{validation_type}/negative_log_likelihood'].update_state( + metrics[f'{validation_type}/negative_log_likelihood'].update_state( # pyrefly: ignore[missing-attribute] negative_log_likelihood ) - metrics[f'{validation_type}/accuracy'].update_state(labels, probs_fm) + metrics[f'{validation_type}/accuracy'].update_state(labels, probs_fm) # pyrefly: ignore[missing-attribute] metrics[f'{validation_type}/ece'].add_batch(probs_fm, label=labels) for _ in tf.range(tf.cast(num_steps, tf.int32)): @@ -750,7 +750,7 @@ def step_fn(inputs): tf.expand_dims(pi_mc_samples, axis=0), multiples=[batch_size, 1, 1, 1]) - _, probs_fm = model((images, random_pi, batched_pi_mc_samples), + _, probs_fm = model((images, random_pi, batched_pi_mc_samples), # pyrefly: ignore[not-callable] training=False) # Remove repeated annotator dimensions: # probs_fm (batch_size, num_dummy_annotators, num_classes) -> @@ -761,9 +761,9 @@ def step_fn(inputs): tf.keras.losses.sparse_categorical_crossentropy(labels, probs_fm)) if dataset_name == 'clean': - metrics[f'{dataset_split}/negative_log_likelihood'].update_state( + metrics[f'{dataset_split}/negative_log_likelihood'].update_state( # pyrefly: ignore[missing-attribute] negative_log_likelihood) - metrics[f'{dataset_split}/accuracy'].update_state(labels, probs_fm) + metrics[f'{dataset_split}/accuracy'].update_state(labels, probs_fm) # pyrefly: ignore[missing-attribute] metrics[f'{dataset_split}/ece'].add_batch(probs_fm, label=labels) elif dataset_name.startswith('ood/'): ood_labels = 1 - inputs['is_in_distribution'] @@ -772,7 +772,7 @@ def step_fn(inputs): # Edgecase for if dataset_name contains underscores for name, metric in metrics.items(): if dataset_name in name: - metric.update_state(ood_labels, ood_scores) + metric.update_state(ood_labels, ood_scores) # pyrefly: ignore[missing-attribute] else: corrupt_metrics['test/nll_{}'.format(dataset_name)].update_state( negative_log_likelihood) @@ -803,7 +803,7 @@ def step_fn(inputs): train_start_time = time.time() train_step(train_iterator) ms_per_example = (time.time() - train_start_time) * 1e6 / batch_size - metrics['train/ms_per_example'].update_state(ms_per_example) + metrics['train/ms_per_example'].update_state(ms_per_example) # pyrefly: ignore[missing-attribute] current_step = (epoch + 1) * steps_per_epoch max_steps = steps_per_epoch * FLAGS.train_epochs @@ -835,7 +835,7 @@ def step_fn(inputs): test_start_time = time.time() test_step(test_iterator, 'test', dataset_name, steps_per_eval) ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size - metrics['test/ms_per_example'].update_state(ms_per_example) + metrics['test/ms_per_example'].update_state(ms_per_example) # pyrefly: ignore[missing-attribute] logging.info('Done with testing on %s', dataset_name) @@ -858,8 +858,8 @@ def step_fn(inputs): logging.info( 'Train Loss: %.4f, Accuracy (FM): %.2f%%, Accuracy (PI): %.2f%%', metrics['train/loss'].result(), - metrics['train/accuracy_fm'].result() * 100, - metrics['train/accuracy_pi'].result() * 100) + metrics['train/accuracy_fm'].result() * 100, # pyrefly: ignore[unsupported-operation] + metrics['train/accuracy_pi'].result() * 100) # pyrefly: ignore[unsupported-operation] logging.info('Test NLL: %.4f, Accuracy: %.2f%%', metrics['test/negative_log_likelihood'].result(), metrics['test/accuracy'].result() * 100) diff --git a/baselines/privileged_information/cifar_pi/tram.py b/baselines/privileged_information/cifar_pi/tram.py index d1b97ea5e..4f5836328 100644 --- a/baselines/privileged_information/cifar_pi/tram.py +++ b/baselines/privileged_information/cifar_pi/tram.py @@ -302,8 +302,8 @@ def main(argv): steps_per_validation = validation_builder.num_examples // batch_size clean_test_builder = ub.datasets.get( 'cifar10' if FLAGS.dataset in ['cifar10n', 'cifar10h'] else 'cifar100', - split=tfds.Split.TEST - if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, + split=tfds.Split.TEST # pyrefly: ignore[missing-attribute] + if FLAGS.dataset != 'cifar10h' else tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] data_dir=data_dir, drop_remainder=FLAGS.drop_remainder_for_eval, is_training=False) @@ -342,7 +342,7 @@ def main(argv): f'{FLAGS.dataset}_corrupted', corruption_type=corruption_type, severity=severity, - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] data_dir=data_dir).load(batch_size=batch_size) test_datasets[f'{corruption_type}_{severity}'] = ( strategy.experimental_distribute_dataset(dataset)) @@ -388,7 +388,7 @@ def annotator_labels_encoding(example): 'random_pi': lambda e: e['pi_features']['random_pi'], } privileged_information_fn = pi_utils.get_privileged_information_fn( - pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) + pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) # pyrefly: ignore[bad-argument-type] pi_shape = privileged_information_fn(dummy_example).shape[1:] @@ -467,12 +467,12 @@ def annotator_labels_encoding(example): rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins), }) if FLAGS.eval_on_ood: - ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names) + ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names) # pyrefly: ignore[unbound-name] metrics.update(ood_metrics) if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, 6): - for corruption in corruption_types: + for corruption in corruption_types: # pyrefly: ignore[unbound-name] dataset_name = '{0}_{1}'.format(corruption, intensity) corrupt_metrics['test/nll_{}'.format(dataset_name)] = ( tf.keras.metrics.Mean()) @@ -545,7 +545,7 @@ def step_fn(inputs): # We take just 1 augmented image from the returned augmented images. images = images[:, 1, ...] with tf.GradientTape() as tape: - (logits, logits_pi) = model((images, privileged_information), + (logits, logits_pi) = model((images, privileged_information), # pyrefly: ignore[not-callable] training=True) # Flatten the annotator axis. @@ -553,7 +553,7 @@ def step_fn(inputs): logits_pi = pi_utils.flatten_annotator_axis(logits_pi) if FLAGS.use_annotator_labels: - logits = tf.gather(logits, non_empty_indices) + logits = tf.gather(logits, non_empty_indices) # pyrefly: ignore[unbound-name] logits_pi = tf.gather(logits_pi, non_empty_indices) if FLAGS.label_smoothing == 0.: @@ -612,7 +612,7 @@ def step_fn(inputs): } for setting in ['pi', 'no_pi']: - metrics[f'train/negative_log_likelihood_{setting}'].update_state( + metrics[f'train/negative_log_likelihood_{setting}'].update_state( # pyrefly: ignore[missing-attribute] stats[setting]['negative_log_likelihood']) for noise_split in ['', '_clean', '_noisy']: sample_weight = None @@ -624,7 +624,7 @@ def step_fn(inputs): stats[setting]['probs'], label=labels, sample_weight=sample_weight) - metrics[f'train/accuracy_{setting}{noise_split}'].update_state( + metrics[f'train/accuracy_{setting}{noise_split}'].update_state( # pyrefly: ignore[missing-attribute] labels, stats[setting]['logits'], sample_weight=sample_weight) metrics['train/loss'].update_state(loss) @@ -664,10 +664,10 @@ def step_fn(inputs): random_pi = privileged_information_fn(dummy_example) random_pi = tf.tile(random_pi, [tf.shape(images)[0], 1, 1]) - logits, _ = model((images, random_pi), training=False) + logits, _ = model((images, random_pi), training=False) # pyrefly: ignore[not-callable] logits = pi_utils.flatten_annotator_axis(logits) if use_annotator_labels: - logits = tf.gather(logits, non_empty_indices) + logits = tf.gather(logits, non_empty_indices) # pyrefly: ignore[unbound-name] probs = tf.nn.softmax(logits) negative_log_likelihood = tf.reduce_mean( @@ -677,9 +677,9 @@ def step_fn(inputs): validation_type = ( 'validation_noisy' if use_annotator_labels else 'validation_clean' ) - metrics[f'{validation_type}/negative_log_likelihood'].update_state( + metrics[f'{validation_type}/negative_log_likelihood'].update_state( # pyrefly: ignore[missing-attribute] negative_log_likelihood) - metrics[f'{validation_type}/accuracy'].update_state(labels, probs) + metrics[f'{validation_type}/accuracy'].update_state(labels, probs) # pyrefly: ignore[missing-attribute] metrics[f'{validation_type}/ece'].add_batch(probs, label=labels) for _ in tf.range(tf.cast(num_steps, tf.int32)): @@ -699,7 +699,7 @@ def step_fn(inputs): random_pi = privileged_information_fn(dummy_example) random_pi = tf.tile(random_pi, [tf.shape(labels)[0], 1, 1]) - logits, _ = model((images, random_pi), training=False) + logits, _ = model((images, random_pi), training=False) # pyrefly: ignore[not-callable] # Remove repeated annotator dimensions: # logits (batch_size, num_dummy_annotators, label_shape) -> # (batch_size, label_shape). @@ -710,9 +710,9 @@ def step_fn(inputs): tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) if dataset_name == 'clean': - metrics[f'{dataset_split}/negative_log_likelihood'].update_state( + metrics[f'{dataset_split}/negative_log_likelihood'].update_state( # pyrefly: ignore[missing-attribute] negative_log_likelihood) - metrics[f'{dataset_split}/accuracy'].update_state(labels, probs) + metrics[f'{dataset_split}/accuracy'].update_state(labels, probs) # pyrefly: ignore[missing-attribute] metrics[f'{dataset_split}/ece'].add_batch(probs, label=labels) elif dataset_name.startswith('ood/'): ood_labels = 1 - inputs['is_in_distribution'] @@ -724,7 +724,7 @@ def step_fn(inputs): # Edgecase for if dataset_name contains underscores for name, metric in metrics.items(): if dataset_name in name: - metric.update_state(ood_labels, ood_scores) + metric.update_state(ood_labels, ood_scores) # pyrefly: ignore[missing-attribute] else: corrupt_metrics['test/nll_{}'.format(dataset_name)].update_state( negative_log_likelihood) @@ -755,7 +755,7 @@ def step_fn(inputs): train_start_time = time.time() train_step(train_iterator) ms_per_example = (time.time() - train_start_time) * 1e6 / batch_size - metrics['train/ms_per_example'].update_state(ms_per_example) + metrics['train/ms_per_example'].update_state(ms_per_example) # pyrefly: ignore[missing-attribute] current_step = (epoch + 1) * steps_per_epoch max_steps = steps_per_epoch * FLAGS.train_epochs @@ -787,31 +787,31 @@ def step_fn(inputs): test_start_time = time.time() test_step(test_iterator, 'test', dataset_name, steps_per_eval) ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size - metrics['test/ms_per_example'].update_state(ms_per_example) + metrics['test/ms_per_example'].update_state(ms_per_example) # pyrefly: ignore[missing-attribute] logging.info('Done with testing on %s', dataset_name) if FLAGS.eval_on_ood: - for ood_dataset_name, ood_dataset in ood_datasets.items(): + for ood_dataset_name, ood_dataset in ood_datasets.items(): # pyrefly: ignore[unbound-name] ood_iterator = iter(ood_dataset) logging.info('Calculating OOD on dataset %s', ood_dataset_name) logging.info('Running OOD eval at epoch: %s', epoch) test_step(ood_iterator, 'test', ood_dataset_name, - steps_per_ood[ood_dataset_name]) + steps_per_ood[ood_dataset_name]) # pyrefly: ignore[unbound-name] logging.info('Done with OOD eval on %s', ood_dataset_name) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): - corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, + corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, # pyrefly: ignore[unbound-name] corruption_types) logging.info( 'Train Loss: %.4f, Accuracy (no PI): %.2f%%, Accuracy (PI): %.2f%%', metrics['train/loss'].result(), - metrics['train/accuracy_no_pi'].result() * 100, - metrics['train/accuracy_pi'].result() * 100) + metrics['train/accuracy_no_pi'].result() * 100, # pyrefly: ignore[unsupported-operation] + metrics['train/accuracy_pi'].result() * 100) # pyrefly: ignore[unsupported-operation] logging.info('Test NLL: %.4f, Accuracy: %.2f%%', metrics['test/negative_log_likelihood'].result(), metrics['test/accuracy'].result() * 100) @@ -831,7 +831,7 @@ def step_fn(inputs): metric.reset_states() if FLAGS.corruptions_interval > 0: - for metric in corrupt_metrics.values(): + for metric in corrupt_metrics.values(): # pyrefly: ignore[unbound-name] metric.reset_states() if (FLAGS.checkpoint_interval > 0 and diff --git a/baselines/privileged_information/imagenet_pi/distill_pi.py b/baselines/privileged_information/imagenet_pi/distill_pi.py index eb1eee9fb..4ad3bab0e 100644 --- a/baselines/privileged_information/imagenet_pi/distill_pi.py +++ b/baselines/privileged_information/imagenet_pi/distill_pi.py @@ -177,7 +177,7 @@ def main(argv): strategy = tf.distribute.TPUStrategy(resolver) train_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TRAIN, + split=tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, one_hot=True, validation_percent=1.0 - FLAGS.train_proportion, @@ -199,7 +199,7 @@ def main(argv): ) steps_per_epoch = train_builder.num_examples // batch_size test_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, data_dir=data_dir, annotations_path=FLAGS.annotations_path, @@ -231,7 +231,7 @@ def main(argv): validation_proportion * 100) # Note we do not one_hot the validation set. validation_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.VALIDATION, + split=tfds.Split.VALIDATION, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, validation_percent=validation_proportion, data_dir=data_dir, @@ -286,7 +286,7 @@ def annotator_label_if_incorrect_encoding(example): annotator_label_if_incorrect_encoding, } privileged_information_fn = pi_utils.get_privileged_information_fn( - pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) + pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) # pyrefly: ignore[bad-argument-type] pi_shape = privileged_information_fn(dummy_example).shape[1:] @@ -313,7 +313,7 @@ def test_annotator_label_if_incorrect_encoding(example): test_privileged_information_fn = pi_utils.get_privileged_information_fn( pi_subset=FLAGS.pi_subset.split(','), - encoding_fn_dict=test_encoding_fn_dict) + encoding_fn_dict=test_encoding_fn_dict) # pyrefly: ignore[bad-argument-type] with strategy.scope(): @@ -487,7 +487,7 @@ def update_train_metrics(labels, labels = tf.argmax(labels, axis=-1) sample_weight = None if noise_split == '_clean': - sample_weight = 1.0 - noisy_idx + sample_weight = 1.0 - noisy_idx # pyrefly: ignore[unsupported-operation] elif noise_split == '_noisy': sample_weight = noisy_idx @@ -677,10 +677,10 @@ def step_fn(inputs): logits = pi_utils.flatten_annotator_axis(logits) if FLAGS.use_bfloat16: - logits = tf.cast(logits, tf.float32) + logits = tf.cast(logits, tf.float32) # pyrefly: ignore[unbound-name] update_test_metrics( - labels, logits, metric_prefix=f'{training_phase}/{metrics_prefix}') + labels, logits, metric_prefix=f'{training_phase}/{metrics_prefix}') # pyrefly: ignore[unbound-name] for _ in tf.range(tf.cast(steps_per_eval, tf.int32)): strategy.run(step_fn, args=(next(iterator),)) @@ -720,7 +720,7 @@ def step_fn(inputs): logging.info('Starting to run %s eval at epoch: %s', training_phase, epoch) if FLAGS.train_proportion < 1.0: - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_clean', iterator=validation_iterator, @@ -728,7 +728,7 @@ def step_fn(inputs): training_phase=training_phase, ) # NOTE: We reinitialize iterator to avoid OutOfRangeError - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_noisy', iterator=validation_iterator, diff --git a/baselines/privileged_information/imagenet_pi/no_pi.py b/baselines/privileged_information/imagenet_pi/no_pi.py index 24fcf3009..2a81f0628 100644 --- a/baselines/privileged_information/imagenet_pi/no_pi.py +++ b/baselines/privileged_information/imagenet_pi/no_pi.py @@ -210,7 +210,7 @@ def main(argv): strategy = tf.distribute.TPUStrategy(resolver) train_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TRAIN, + split=tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, one_hot=True, validation_percent=1.0 - FLAGS.train_proportion, @@ -231,7 +231,7 @@ def main(argv): ) steps_per_epoch = train_builder.num_examples // batch_size test_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, data_dir=data_dir, annotations_path=FLAGS.annotations_path, @@ -258,7 +258,7 @@ def main(argv): validation_proportion * 100) # Note we do not one_hot the validation set. validation_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.VALIDATION, + split=tfds.Split.VALIDATION, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, validation_percent=validation_proportion, data_dir=data_dir, @@ -445,7 +445,7 @@ def step_fn(inputs): metrics[f'train/ece{noise_split}'].add_batch( probs, label=labels, sample_weight=sample_weight) metrics[f'train/accuracy{noise_split}'].update_state( - labels, logits, sample_weight=sample_weight) + labels, logits, sample_weight=sample_weight) # pyrefly: ignore[bad-keyword-argument] metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) @@ -519,13 +519,13 @@ def step_fn(inputs): test_iterator = iter(test_dataset) logging.info('Starting to run eval at epoch: %s', epoch) if FLAGS.train_proportion < 1.0: - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_clean', iterator=validation_iterator, steps_per_eval=steps_per_validation_eval) # NOTE: We reinitialize the iterator to avoid OutOfRangeError. - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_noisy', iterator=validation_iterator, diff --git a/baselines/privileged_information/imagenet_pi/pi_full_marginalization.py b/baselines/privileged_information/imagenet_pi/pi_full_marginalization.py index d1886d86b..10e4c5fa2 100644 --- a/baselines/privileged_information/imagenet_pi/pi_full_marginalization.py +++ b/baselines/privileged_information/imagenet_pi/pi_full_marginalization.py @@ -208,7 +208,7 @@ def main(argv): strategy = tf.distribute.TPUStrategy(resolver) train_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TRAIN, + split=tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, one_hot=True, validation_percent=1.0 - FLAGS.train_proportion, @@ -230,7 +230,7 @@ def main(argv): ) steps_per_epoch = train_builder.num_examples // batch_size test_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, data_dir=data_dir, annotations_path=FLAGS.annotations_path, @@ -262,7 +262,7 @@ def main(argv): validation_proportion * 100) # Note we do not one_hot the validation set. validation_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.VALIDATION, + split=tfds.Split.VALIDATION, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, validation_percent=validation_proportion, data_dir=data_dir, @@ -317,7 +317,7 @@ def annotator_label_if_incorrect_encoding(example): annotator_label_if_incorrect_encoding, } privileged_information_fn = pi_utils.get_privileged_information_fn( - pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) + pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) # pyrefly: ignore[bad-argument-type] pi_shape = privileged_information_fn(dummy_example).shape[1:] pi_mc_samples = pi_utils.sample_pi_features(train_builder, @@ -347,7 +347,7 @@ def test_annotator_label_if_incorrect_encoding(example): test_privileged_information_fn = pi_utils.get_privileged_information_fn( pi_subset=FLAGS.pi_subset.split(','), - encoding_fn_dict=test_encoding_fn_dict) + encoding_fn_dict=test_encoding_fn_dict) # pyrefly: ignore[bad-argument-type] with strategy.scope(): @@ -540,7 +540,7 @@ def step_fn(inputs): label=labels, sample_weight=sample_weight) metrics[f'train/accuracy_{setting}{noise_split}'].update_state( - labels, stats[setting]['outputs'], sample_weight=sample_weight) # pylint: disable=redundant-keyword-arg + labels, stats[setting]['outputs'], sample_weight=sample_weight) # pylint: disable=redundant-keyword-arg # pyrefly: ignore[bad-keyword-argument] metrics['train/loss'].update_state(loss) @@ -639,13 +639,13 @@ def step_fn(inputs): test_iterator = iter(test_dataset) logging.info('Starting to run eval at epoch: %s', epoch) if FLAGS.train_proportion < 1.0: - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_noisy', iterator=validation_iterator, steps_per_eval=steps_per_validation_eval) # NOTE: We reinitialize the iterator to avoid OutOfRangeError. - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_clean', iterator=validation_iterator, diff --git a/baselines/privileged_information/imagenet_pi/tram.py b/baselines/privileged_information/imagenet_pi/tram.py index c9b143850..7f09c3fd3 100644 --- a/baselines/privileged_information/imagenet_pi/tram.py +++ b/baselines/privileged_information/imagenet_pi/tram.py @@ -226,7 +226,7 @@ def main(argv): strategy = tf.distribute.TPUStrategy(resolver) train_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TRAIN, + split=tfds.Split.TRAIN, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, one_hot=True, validation_percent=1.0 - FLAGS.train_proportion, @@ -248,7 +248,7 @@ def main(argv): ) steps_per_epoch = train_builder.num_examples // batch_size test_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.TEST, + split=tfds.Split.TEST, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, data_dir=data_dir, annotations_path=FLAGS.annotations_path, @@ -280,7 +280,7 @@ def main(argv): validation_proportion * 100) # Note we do not one_hot the validation set. validation_builder = ub.datasets.ImageNetPIDataset( - split=tfds.Split.VALIDATION, + split=tfds.Split.VALIDATION, # pyrefly: ignore[missing-attribute] use_bfloat16=FLAGS.use_bfloat16, validation_percent=validation_proportion, data_dir=data_dir, @@ -335,7 +335,7 @@ def annotator_label_if_incorrect_encoding(example): annotator_label_if_incorrect_encoding, } privileged_information_fn = pi_utils.get_privileged_information_fn( - pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) + pi_subset=FLAGS.pi_subset.split(','), encoding_fn_dict=encoding_fn_dict) # pyrefly: ignore[bad-argument-type] pi_shape = privileged_information_fn(dummy_example).shape[1:] @@ -360,7 +360,7 @@ def test_annotator_label_if_incorrect_encoding(example): test_privileged_information_fn = pi_utils.get_privileged_information_fn( pi_subset=FLAGS.pi_subset.split(','), - encoding_fn_dict=test_encoding_fn_dict) + encoding_fn_dict=test_encoding_fn_dict) # pyrefly: ignore[bad-argument-type] with strategy.scope(): @@ -557,7 +557,7 @@ def step_fn(inputs): label=labels, sample_weight=sample_weight) metrics[f'train/accuracy_{setting}{noise_split}'].update_state( - labels, stats[setting]['logits'], sample_weight=sample_weight) + labels, stats[setting]['logits'], sample_weight=sample_weight) # pyrefly: ignore[bad-keyword-argument] metrics['train/loss'].update_state(loss) @@ -637,13 +637,13 @@ def step_fn(inputs): test_iterator = iter(test_dataset) logging.info('Starting to run eval at epoch: %s', epoch) if FLAGS.train_proportion < 1.0: - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_noisy', iterator=validation_iterator, steps_per_eval=steps_per_validation_eval) # NOTE: We reinitialize the iterator to avoid OutOfRangeError. - validation_iterator = iter(validation_dataset) + validation_iterator = iter(validation_dataset) # pyrefly: ignore[no-matching-overload] test_step( metrics_prefix='validation_clean', iterator=validation_iterator, diff --git a/baselines/t5/data/deepbank/graph_linear_utils.py b/baselines/t5/data/deepbank/graph_linear_utils.py index 08eaebce3..745716f8e 100644 --- a/baselines/t5/data/deepbank/graph_linear_utils.py +++ b/baselines/t5/data/deepbank/graph_linear_utils.py @@ -109,16 +109,16 @@ def _configure(triples: Triples, root: str) -> tree_utils.Tree: # 'x2': ('x2', [('/', '_udef_q')])} # Initialization of nodemap. nodemap = {idx: None for idx in node_indexes} - nodemap[root] = (root, []) + nodemap[root] = (root, []) # pyrefly: ignore[unsupported-operation] data = list(reversed(triples)) # Initialization of node. - node = _configure_node(root, data, nodemap) + node = _configure_node(root, data, nodemap) # pyrefly: ignore[bad-argument-type] while data: - skipped, idx, data = _find_next(data, nodemap) + skipped, idx, data = _find_next(data, nodemap) # pyrefly: ignore[bad-argument-type] data_count = len(data) if idx is None or data_count == 0: raise ValueError('Possibly disconnected graph') - _configure_node(idx, data, nodemap) + _configure_node(idx, data, nodemap) # pyrefly: ignore[bad-argument-type] if len(data) >= data_count: raise ValueError('Possible cycle in configuration') data = skipped + data diff --git a/baselines/t5/data/deepbank/graph_utils.py b/baselines/t5/data/deepbank/graph_utils.py index f073c74d2..a587b2a0e 100644 --- a/baselines/t5/data/deepbank/graph_utils.py +++ b/baselines/t5/data/deepbank/graph_utils.py @@ -941,10 +941,10 @@ def get_dag_match(cur_dag1, logging.warning("Fail to parse DAG: %s", cur_dag2) dag2_parsed = False - if dag1_parsed and not dag1.root: + if dag1_parsed and not dag1.root: # pyrefly: ignore[unbound-name] logging.warning("Fail to parse DAG: %s", cur_dag1) dag1_parsed = False - if dag2_parsed and not dag2.root: + if dag2_parsed and not dag2.root: # pyrefly: ignore[unbound-name] logging.warning("Fail to parse DAG: %s", cur_dag2) dag2_parsed = False @@ -967,8 +967,8 @@ def get_dag_match(cur_dag1, compare_instance = compare_attribute = False if dag1_parsed and dag2_parsed: (best_mapping, best_match_num) = get_best_match( - instance1, attribute1, relation1, - instance2, attribute2, relation2, + instance1, attribute1, relation1, # pyrefly: ignore[unbound-name] + instance2, attribute2, relation2, # pyrefly: ignore[unbound-name] prefix1, prefix2, compare_instance=compare_instance, compare_attribute=compare_attribute, diff --git a/baselines/t5/data/deepbank/lispress_utils.py b/baselines/t5/data/deepbank/lispress_utils.py index 7ac2cb5cc..b5baff9a6 100644 --- a/baselines/t5/data/deepbank/lispress_utils.py +++ b/baselines/t5/data/deepbank/lispress_utils.py @@ -163,7 +163,7 @@ def lispress_to_graph(sexp: Sexp) -> str: if len(sexp) == 1: # Input example: ['Constraint[Recipient]']. # Output example: ' ( Constraint[Recipient] ) '. - return ' ( ' + sexp[0] + ' ) ' + return ' ( ' + sexp[0] + ' ) ' # pyrefly: ignore[unsupported-operation] elif len(sexp) == 2: if isinstance(sexp[1], str): if sexp[0] == 'Number': diff --git a/baselines/t5/data/deepbank/penman_utils.py b/baselines/t5/data/deepbank/penman_utils.py index 44f9c718f..61c866d40 100644 --- a/baselines/t5/data/deepbank/penman_utils.py +++ b/baselines/t5/data/deepbank/penman_utils.py @@ -259,7 +259,7 @@ def retoken_graph_str(graph_str: Text, data_version: str = 'v0') -> Text: edge_name = token.split('-of')[0] if edge_name in name_to_token_map: retoken = token.replace(edge_name, name_to_token_map[edge_name]) - if '-of' in retoken and '-of' in name_to_token_map: + if '-of' in retoken and '-of' in name_to_token_map: # pyrefly: ignore[unbound-name] retoken = retoken.replace('-of', ' ' + name_to_token_map['-of']) new_graph_str_list.append(retoken) elif token.startswith('_') and data_version in ['v0', 'v1']: diff --git a/baselines/t5/data/tasks/deepbank.py b/baselines/t5/data/tasks/deepbank.py index 77c04eb53..272083ef4 100644 --- a/baselines/t5/data/tasks/deepbank.py +++ b/baselines/t5/data/tasks/deepbank.py @@ -152,12 +152,12 @@ def parsing_dataset( datasets = [ tfrecord_dataset( data_patterns=[p], shuffle_files=shuffle_files, repeat=True) - for p in params['train_patterns'] + for p in params['train_patterns'] # pyrefly: ignore[unsupported-operation] ] - return dataset_mixture(datasets, weights=params['train_weights']) + return dataset_mixture(datasets, weights=params['train_weights']) # pyrefly: ignore[unsupported-operation] return tfrecord_dataset( - data_patterns=[params['eval_patterns'][split]], + data_patterns=[params['eval_patterns'][split]], # pyrefly: ignore[unsupported-operation] shuffle_files=shuffle_files, repeat=False) diff --git a/baselines/t5/data/tasks/mnli.py b/baselines/t5/data/tasks/mnli.py index e0738dfd7..ac8829a2d 100644 --- a/baselines/t5/data/tasks/mnli.py +++ b/baselines/t5/data/tasks/mnli.py @@ -151,7 +151,7 @@ def register_hans_dataset( TaskRegistry.add( data_name, - source=data_source, + source=data_source, # pyrefly: ignore[bad-argument-type] preprocessors=[ hans_tsv_preprocessor, hans_example_preprocessor, @@ -182,7 +182,7 @@ def register_mnli_dataset( TaskRegistry.add( data_name, - source=tfds_source, + source=tfds_source, # pyrefly: ignore[bad-argument-type] preprocessors=[ mnli_preprocessor, seqio.preprocessors.tokenize, diff --git a/baselines/t5/data/tasks/nalue.py b/baselines/t5/data/tasks/nalue.py index 0ad5a1c13..cc0ea658c 100644 --- a/baselines/t5/data/tasks/nalue.py +++ b/baselines/t5/data/tasks/nalue.py @@ -116,7 +116,7 @@ def register_nalue_dataset(data_name: str, split_names: Sequence[str], file_num_examples: Mapping[str, int]): """Registers NaLUE datasets.""" # Filters file metadata by split_names. - file_patterns = {k: v for k, v in file_patterns.items() if k in split_names} + file_patterns = {k: v for k, v in file_patterns.items() if k in split_names} # pyrefly: ignore[bad-assignment] file_num_examples = { k: v for k, v in file_num_examples.items() if k in split_names } @@ -160,7 +160,7 @@ def register_nalue_dataset(data_name: str, split_names: Sequence[str], TaskRegistry.add( data_name, - source=data_source, + source=data_source, # pyrefly: ignore[bad-argument-type] preprocessors=[ functools.partial( t5_preprocessors.parse_tsv, field_names=NALUE_TSV_FIELD_NAMES), @@ -208,11 +208,11 @@ def register_nalue_dataset(data_name: str, split_names: Sequence[str], register_nalue_dataset( 'nalue_ind_and_standard_oos', split_names=('validation', 'test'), - file_patterns=NALUE_FILE_PATTERNS_IND_AND_OOS, + file_patterns=NALUE_FILE_PATTERNS_IND_AND_OOS, # pyrefly: ignore[bad-argument-type] file_num_examples=NALUE_NUM_EXAMPLES_IND_AND_OOS) register_nalue_dataset( 'nalue_ind_and_near_oos', split_names=('validation', 'test'), - file_patterns=NALUE_FILE_PATTERNS_IND_AND_NEAR_OOS, + file_patterns=NALUE_FILE_PATTERNS_IND_AND_NEAR_OOS, # pyrefly: ignore[bad-argument-type] file_num_examples=NALUE_NUM_EXAMPLES_IND_AND_NEAR_OOS) diff --git a/baselines/t5/data/tasks/toxic_comments.py b/baselines/t5/data/tasks/toxic_comments.py index 2abe70e20..a8a28e65b 100644 --- a/baselines/t5/data/tasks/toxic_comments.py +++ b/baselines/t5/data/tasks/toxic_comments.py @@ -81,7 +81,7 @@ def _register_toxic_comments_ranking_task( TaskRegistry.add( task_name, - source=tfds_source, + source=tfds_source, # pyrefly: ignore[bad-argument-type] preprocessors=[ toxicity_rank_classification_preprocessor, toxicity_rank_classification_formatter, seqio.preprocessors.tokenize, @@ -126,7 +126,7 @@ def _register_toxic_comments_classification_task( # outputs. TaskRegistry.add( task_name, - source=tfds_source, + source=tfds_source, # pyrefly: ignore[bad-argument-type] preprocessors=[ toxicity_classification_preprocessor, seqio.preprocessors.tokenize, diff --git a/baselines/t5/models/be_models.py b/baselines/t5/models/be_models.py index 7d1aea333..ad5ba7f13 100644 --- a/baselines/t5/models/be_models.py +++ b/baselines/t5/models/be_models.py @@ -62,7 +62,7 @@ def _compute_argmax_score( if return_intermediates: sequence_scores, intermediates = sequence_scores if ens_size > 1: - sequence_scores = jnp.reshape(sequence_scores, + sequence_scores = jnp.reshape(sequence_scores, # pyrefly: ignore[bad-argument-type] (ens_size, -1) + sequence_scores.shape[1:]) # pytype: disable=attribute-error # jax-ndarray if ensemble_probs: # Computes log(mean(exp(logits))) along the first dimension. @@ -71,7 +71,7 @@ def _compute_argmax_score( else: sequence_scores = jnp.mean(sequence_scores, axis=0) if return_intermediates: - return sequence_scores, intermediates + return sequence_scores, intermediates # pyrefly: ignore[bad-return, unbound-name] return sequence_scores def _compute_logits_from_slice(self, diff --git a/baselines/t5/models/gp_models.py b/baselines/t5/models/gp_models.py index f4c94a062..aa759db6a 100644 --- a/baselines/t5/models/gp_models.py +++ b/baselines/t5/models/gp_models.py @@ -53,7 +53,7 @@ def loss_fn( weights=weights, label_smoothing=self._label_smoothing, z_loss=self._z_loss, - loss_normalizing_factor=self._loss_normalizing_factor) + loss_normalizing_factor=self._loss_normalizing_factor) # pyrefly: ignore[bad-argument-type] metrics = models.compute_base_metrics(logits, targets, weights, loss, total_z_loss) diff --git a/baselines/t5/models/heteroscedastic_models.py b/baselines/t5/models/heteroscedastic_models.py index 2be609590..194056ec5 100644 --- a/baselines/t5/models/heteroscedastic_models.py +++ b/baselines/t5/models/heteroscedastic_models.py @@ -127,4 +127,4 @@ def get_assignment_map( # `scale_layer_heteroscedastic` in the parameter-efficient case. (r'(.*)decoder/heteroscedastic_head/scale_layer(.*)', None), ] - return assignment_map + return assignment_map # pyrefly: ignore[bad-return] diff --git a/baselines/t5/models/models.py b/baselines/t5/models/models.py index 28d45633b..1faf7c070 100644 --- a/baselines/t5/models/models.py +++ b/baselines/t5/models/models.py @@ -142,7 +142,7 @@ def _get_label_token_ids(self): label_token_ids = [] for token in self.label_tokens: - token_id = self._output_vocabulary.encode(token) + token_id = self._output_vocabulary.encode(token) # pyrefly: ignore[bad-argument-type] # Convert token_id to integer. if not isinstance(token_id, INT_CLASSES): @@ -235,7 +235,7 @@ def get_sequence_scores_and_intermediates(vals): intermediates = None # Returns class logits, shape (batch_size, output_len, num_class). - sequence_scores = logits[:, :, self.label_token_ids] + sequence_scores = logits[:, :, self.label_token_ids] # pyrefly: ignore[bad-index] return sequence_scores, intermediates batch_sequence_scores, intermediates = _partial_map( @@ -300,7 +300,7 @@ def score_batch( if dropout_seed is not None: dropout_rng = jax.random.PRNGKey(dropout_seed) if num_mcdropout_samples is not None: - dropout_rng = jax.random.split(dropout_rng, num_mcdropout_samples) + dropout_rng = jax.random.split(dropout_rng, num_mcdropout_samples) # pyrefly: ignore[bad-argument-type] # Computes argmax predictive logits over the whole class, # shape (batch_size, output_len, num_class). sequence_scores = self._compute_argmax_score(params, batch, @@ -339,7 +339,7 @@ def score_batch( rngs=dropout_rng, method=self.module.encode, ) - intermediates['encoded_inputs'] = encoded_inputs + intermediates['encoded_inputs'] = encoded_inputs # pyrefly: ignore[unbound-name, unsupported-operation] intermediates_flat = traverse_util.flatten_dict(intermediates) intermediates = {} @@ -347,7 +347,7 @@ def score_batch( key = '/'.join(key_flat) if key in intermediates_to_track: intermediates[key] = value - return sequence_scores, intermediates + return sequence_scores, intermediates # pyrefly: ignore[unbound-name] return sequence_scores @@ -455,7 +455,7 @@ def get_flat_logits_and_new_cache(vals): new_cache = _cache_concatenate(batch_new_cache) return flat_logits, new_cache - def predict_batch_with_aux( + def predict_batch_with_aux( # pyrefly: ignore[bad-override] self, params: PyTree, batch: Mapping[str, jnp.ndarray], @@ -502,7 +502,7 @@ def predict_batch_with_aux( elif run_mc_dropout: params_list = [params] batch_rngs = {'dropout': dropout_rng} - k = dropout_rng.shape[0] + k = dropout_rng.shape[0] # pyrefly: ignore[missing-attribute] else: params_list = [params] k = 1 @@ -541,7 +541,7 @@ def get_cache(args): def get_encoded_inputs(args): params, rngs = args encoded_inputs = decoding.flat_batch_beam_expand( - self.module.apply({'params': params}, + self.module.apply({'params': params}, # pyrefly: ignore[bad-argument-type] inputs, enable_dropout=rngs is not None, rngs=rngs, @@ -587,7 +587,7 @@ def get_encoded_inputs(args): # decodes: [k * batch, num_decodes, max_decode_len + 1] # scores: [k * batch, num_decodes] scanned = hasattr(self.module, 'scan_layers') and self.module.scan_layers - decodes, scores = self._decode_fn( + decodes, scores = self._decode_fn( # pyrefly: ignore[not-callable] inputs=decoder_prompt_inputs, cache=cache, tokens_to_logits=tokens_ids_to_logits, @@ -626,7 +626,7 @@ class EncoderDecoderBeamScoreModel(EncoderDecoderClassifierModel): `return_all_decodes` set to `True`. """ - def predict_batch( + def predict_batch( # pyrefly: ignore[bad-override] self, params: PyTree, batch: Mapping[str, jnp.ndarray], @@ -704,7 +704,7 @@ def predict_batch_with_aux( # pytype: disable=signature-mismatch # overriding- if dropout_seed is not None: dropout_rng = jax.random.PRNGKey(dropout_seed) if num_mcdropout_samples is not None: - dropout_rng = jax.random.split(dropout_rng, num_mcdropout_samples) + dropout_rng = jax.random.split(dropout_rng, num_mcdropout_samples) # pyrefly: ignore[bad-argument-type] return super().predict_batch_with_aux( params=params, batch=batch, @@ -747,7 +747,7 @@ def score_batch( # pytype: disable=signature-mismatch # overriding-parameter-c if dropout_seed is not None: dropout_rng = jax.random.PRNGKey(dropout_seed) if num_mcdropout_samples is not None: - dropout_rng = jax.random.split(dropout_rng, num_mcdropout_samples) + dropout_rng = jax.random.split(dropout_rng, num_mcdropout_samples) # pyrefly: ignore[bad-argument-type] # Computes argmax predictive logits over the whole class, # shape (batch_size, output_len, num_class). sequence_scores = self._compute_argmax_score(params, batch, @@ -767,15 +767,15 @@ def score_batch( # pytype: disable=signature-mismatch # overriding-parameter-c token_scores = -losses.cross_entropy_with_logits( logits, common_utils.onehot( - target_tokens, jnp.shape(logits)[-1], on_value=1, off_value=0), + target_tokens, jnp.shape(logits)[-1], on_value=1, off_value=0), # pyrefly: ignore[bad-argument-type] z_loss=0.0)[0] * weights sequence_scores = token_scores.sum(-1) if return_intermediates: - intermediates.setdefault('entropy', {}) + intermediates.setdefault('entropy', {}) # pyrefly: ignore[missing-attribute, unbound-name] intermediates['entropy']['logits'] = logits - intermediates['entropy']['token_entropy'] = _compute_token_entropy(logits) + intermediates['entropy']['token_entropy'] = _compute_token_entropy(logits) # pyrefly: ignore[bad-argument-type] if intermediates_to_track: # infer binary's write_fn requires intermediates to be a flatten dict. # So we will return flatten dictionary here. For example, given @@ -809,7 +809,7 @@ class DecoderOnlyBeamScoreModel(t5x_models.DecoderOnlyModel): decoding algorithms. """ - def predict_batch( + def predict_batch( # pyrefly: ignore[bad-override] self, params: PyTree, batch: Mapping[str, jnp.ndarray], diff --git a/experimental/language_structure/psl/data.py b/experimental/language_structure/psl/data.py index d4091f1d7..4f575223e 100644 --- a/experimental/language_structure/psl/data.py +++ b/experimental/language_structure/psl/data.py @@ -221,10 +221,10 @@ def create_utterance_mask_feature(dialogs: tf.Tensor, """Creates features from actual dialog turn length.""" utterance_mask = _get_utterance_mask(dialogs) last_utterance_mask = _get_last_utterance_mask(utterance_mask) - non_last_utterance_mask = utterance_mask - last_utterance_mask - pad_utterance_mask = 1 - utterance_mask + non_last_utterance_mask = utterance_mask - last_utterance_mask # pyrefly: ignore[unsupported-operation] + pad_utterance_mask = 1 - utterance_mask # pyrefly: ignore[unsupported-operation] mask_feature = ( - last_utterance_mask * last_utterance_mask_value + + last_utterance_mask * last_utterance_mask_value + # pyrefly: ignore[unsupported-operation] non_last_utterance_mask * utterance_mask_value + pad_utterance_mask * pad_utterance_mask_value) return mask_feature @@ -302,7 +302,7 @@ def create_features(user_utterance_ids: tf.Tensor, ] for i, keyword_ids in enumerate(keyword_ids_per_class): if check_keyword_by_utterance: - dialogs = candidate_dialogs[i % 2] + dialogs = candidate_dialogs[i % 2] # pyrefly: ignore[unbound-name] features.append( create_keyword_feature(dialogs, keyword_ids, include_keyword_value, exclude_keyword_value)) @@ -399,7 +399,7 @@ def pad_one_hot_labels( def list_to_dataset(data: List[Dialog], labels: List[List[List[int]]], shuffle: bool, batch_size: int) -> tf.data.Dataset: """Converts list into tensorflow dataset.""" - ds = tf.data.Dataset.from_tensor_slices((data, labels)) + ds = tf.data.Dataset.from_tensor_slices((data, labels)) # pyrefly: ignore[bad-argument-type] if shuffle: ds = ds.shuffle(buffer_size=len(data)) ds = ds.batch(batch_size) diff --git a/experimental/language_structure/psl/psl_model.py b/experimental/language_structure/psl/psl_model.py index fda88bf75..c1211d7d2 100644 --- a/experimental/language_structure/psl/psl_model.py +++ b/experimental/language_structure/psl/psl_model.py @@ -82,7 +82,7 @@ def compute_rule_loss(body: List[tf.Tensor], return tf.reduce_mean(-tf.reduce_sum(tf.math.log(prob), axis=-1)) if logic == 'product_real_logic': - body = PSLModel.soft_and(body, logic) + body = PSLModel.soft_and(body, logic) # pyrefly: ignore[bad-assignment] truth_values = PSLModel.soft_imply(body, head, logic) if loss_function == _LINEAR_LOSS: @@ -103,10 +103,10 @@ def soft_imply( """Soft logical implication.""" if logic == 'lukasiewicz': # body -> head = (1 - body_i), head - return [1. - predicate for predicate in body], head + return [1. - predicate for predicate in body], head # pyrefly: ignore[bad-return, unsupported-operation] if logic == 'product_real_logic': # 1 - A + A * B - return 1 - body + body * head + return 1 - body + body * head # pyrefly: ignore[bad-return, unsupported-operation] else: raise ValueError('Unsuported logic: %s' % (logic,)) @@ -116,9 +116,9 @@ def soft_imply2(body: List[tf.Tensor], logic: str = 'lukasiewicz') -> tf.Tensor: """Soft logical implication.""" if logic == 'lukasiewicz': - body = PSLModel.soft_and(body) + body = PSLModel.soft_and(body) # pyrefly: ignore[bad-assignment] # A -> B = min(1, 1 - A + B) - return tf.minimum(1., 1 - body + head) + return tf.minimum(1., 1 - body + head) # pyrefly: ignore[unsupported-operation] else: raise ValueError('Unsuported logic: %s' % (logic,)) @@ -132,7 +132,7 @@ def soft_and(body: List[tf.Tensor], logic: str = _DEFAULT_LOGIC) -> tf.Tensor: # A & B = max(0, A + B - 1) val = body[0] for i in range(1, num_bodies): - val = normalize(tf.nn.relu(val + body[i] - 1)) + val = normalize(tf.nn.relu(val + body[i] - 1)) # pyrefly: ignore[unsupported-operation] return val elif logic == 'product_real_logic': # A & B = A * B @@ -148,9 +148,9 @@ def soft_not(predicate: tf.Tensor, logic: str = _DEFAULT_LOGIC) -> tf.Tensor: """Soft logical negation.""" if logic == 'lukasiewicz': # !predicate = 1 - predicate - return normalize(1. - predicate) + return normalize(1. - predicate) # pyrefly: ignore[unsupported-operation] if logic == 'product_real_logic': - return 1. - predicate + return 1. - predicate # pyrefly: ignore[bad-return, unsupported-operation] else: raise ValueError('Unsuported logic: %s' % (logic,)) diff --git a/experimental/language_structure/psl/psl_model_dstc_synthetic.py b/experimental/language_structure/psl/psl_model_dstc_synthetic.py index fdb2b8fcd..c872680e1 100644 --- a/experimental/language_structure/psl/psl_model_dstc_synthetic.py +++ b/experimental/language_structure/psl/psl_model_dstc_synthetic.py @@ -129,7 +129,7 @@ def _previous_statement(self, data): def generate_predicates(self, input_ids: Sequence[tf.Tensor]): """Generates potentials used throughout the rules.""" - hidden = self.embed_layer(input_ids) + hidden = self.embed_layer(input_ids) # pyrefly: ignore[not-callable] input_mask = tf.sign(input_ids) logits = tf.reduce_sum( hidden * diff --git a/experimental/language_structure/psl/psl_model_multiwoz_test_util.py b/experimental/language_structure/psl/psl_model_multiwoz_test_util.py index 9e899f85c..225c880d8 100644 --- a/experimental/language_structure/psl/psl_model_multiwoz_test_util.py +++ b/experimental/language_structure/psl/psl_model_multiwoz_test_util.py @@ -218,11 +218,11 @@ def build_constrained_model(input_size: List[int]) -> tf.keras.Model: """Build simple neural model for class prediction.""" input_layer = tf.keras.layers.Input(input_size) - hidden_layer_1 = tf.keras.layers.Dense(1024)(input_layer) - hidden_layer_2 = tf.keras.layers.Dense( + hidden_layer_1 = tf.keras.layers.Dense(1024)(input_layer) # pyrefly: ignore[not-callable] + hidden_layer_2 = tf.keras.layers.Dense( # pyrefly: ignore[not-callable] 512, activation='sigmoid')( hidden_layer_1) - output = tf.keras.layers.Dense( + output = tf.keras.layers.Dense( # pyrefly: ignore[not-callable] 9, activation='softmax', kernel_regularizer=tf.keras.regularizers.l2(1.0))( hidden_layer_2) diff --git a/experimental/multimodal/checkpoint_utils.py b/experimental/multimodal/checkpoint_utils.py index 5bb5cf6ac..5a682a2d2 100644 --- a/experimental/multimodal/checkpoint_utils.py +++ b/experimental/multimodal/checkpoint_utils.py @@ -123,7 +123,7 @@ def read_chunk(i): # Fill in the empty `data` array in parallel. pool = thread.ThreadPoolExecutor(pool_size) - results = pool.map(read_chunk, range(int(num_bufs) + 1)) + results = pool.map(read_chunk, range(int(num_bufs) + 1)) # pyrefly: ignore[unbound-name] pool.shutdown(wait=False) logging.debug("results: %s", list(results)) return data @@ -168,7 +168,7 @@ def load_checkpoint(tree: Optional[Params], treedef = jax.tree_util.tree_structure(tree) tree = jax.tree_util.tree_unflatten(treedef, values) else: - tree = _recover_tree(keys, values) + tree = _recover_tree(keys, values) # pyrefly: ignore[unbound-name] return tree @@ -267,7 +267,7 @@ def checkpoint_trained_model( accum_train_time=checkpoint_data.accumulated_train_time), ) if checkpoint_data.fixed_model_states is not None: - tree["states"] = checkpoint_data.fixed_model_states + tree["states"] = checkpoint_data.fixed_model_states # pyrefly: ignore[bad-assignment] save_checkpoint(tree, path, step_for_copy) @@ -278,7 +278,7 @@ def _flatten_jax_params_dict(d: Params, parent_key: str = "", for k, v in d.items(): path = parent_key + sep + k if parent_key else k if isinstance(v, collections.abc.Mapping): - items.extend(_flatten_jax_params_dict(v, path, sep=sep).items()) + items.extend(_flatten_jax_params_dict(v, path, sep=sep).items()) # pyrefly: ignore[bad-argument-type] else: items.append((path, v)) @@ -527,7 +527,7 @@ def maybe_load_checkpoint(train_loop_rngs: jnp.ndarray, checkpoint_tree = {"opt": init_optimizer, "extra": checkpoint_extra} if init_fixed_model_states is not None: - checkpoint_tree["states"] = init_fixed_model_states + checkpoint_tree["states"] = init_fixed_model_states # pyrefly: ignore[bad-assignment] checkpoint = load_checkpoint(checkpoint_tree, resume_checkpoint_path) optimizer, checkpoint_extra = checkpoint["opt"], checkpoint["extra"] fixed_model_states = checkpoint.get("states", {}) @@ -566,7 +566,7 @@ def maybe_load_checkpoint(train_loop_rngs: jnp.ndarray, return CheckpointData( optimizer=optimizer, fixed_model_states=fixed_model_states, - train_loop_rngs=checkpoint_extra["rngs_loop"], + train_loop_rngs=checkpoint_extra["rngs_loop"], # pyrefly: ignore[bad-argument-type] accumulated_train_time=checkpoint_extra["accum_train_time"]) diff --git a/experimental/multimodal/deterministic.py b/experimental/multimodal/deterministic.py index 9ab39be7c..a2b1c3905 100644 --- a/experimental/multimodal/deterministic.py +++ b/experimental/multimodal/deterministic.py @@ -127,7 +127,7 @@ def write_note(note): split=config.train_split, rng=train_ds_rng, process_batch_size=local_batch_size, - preprocess_fn=preprocess_spec.parse( + preprocess_fn=preprocess_spec.parse( # pyrefly: ignore[bad-argument-type] spec=config.pp_train, available_ops=preprocess_utils.all_ops()), shuffle_buffer_size=config.shuffle_buffer_size, prefetch_size=config.get('prefetch_to_host', 2), @@ -392,7 +392,7 @@ def loss_fn(params, states, images, texts): # Optionally resize the global gradient to a maximum norm. We found this # useful in some cases across optimizers, hence it's in the main loop. if config.get('grad_clip_norm'): - g_factor = jnp.minimum(1.0, config.grad_clip_norm / l2_g) + g_factor = jnp.minimum(1.0, config.grad_clip_norm / l2_g) # pyrefly: ignore[unbound-name] g = jax.tree_util.tree_map(lambda p: g_factor * p, g) opt = opt.apply_gradient(g, learning_rate=lr) @@ -411,7 +411,7 @@ def loss_fn(params, states, images, texts): reint_params = [] checkpoint_data = checkpoint_utils.maybe_load_checkpoint( train_loop_rngs=rng, - save_checkpoint_path=save_checkpoint_path, + save_checkpoint_path=save_checkpoint_path, # pyrefly: ignore[bad-argument-type] init_optimizer=opt_cpu, init_params=params_cpu, init_fixed_model_states=states_cpu, @@ -533,7 +533,7 @@ def loss_fn(params, states, images, texts): timing_measurements, note = chrono.tick(step) write_note(note) train_measurements = {} - train_measurements.update(flax.jax_utils.unreplicate(extra_measurements)) + train_measurements.update(flax.jax_utils.unreplicate(extra_measurements)) # pyrefly: ignore[unbound-name] train_measurements.update(timing_measurements) writer.write_scalars(step, train_measurements) # Keep train_loss to return for reproducibility tests. diff --git a/experimental/multimodal/input_utils.py b/experimental/multimodal/input_utils.py index 6c61aa7d5..92a587e39 100644 --- a/experimental/multimodal/input_utils.py +++ b/experimental/multimodal/input_utils.py @@ -166,7 +166,7 @@ def _build_dataset(dataset: Union[str, tfds.core.DatasetBuilder], dataset_options.experimental_threading.max_intra_op_parallelism = 1 read_config = tfds.ReadConfig( # pytype: disable=wrong-arg-types # jax-ndarray - shuffle_seed=file_shuffle_seed, options=dataset_options) + shuffle_seed=file_shuffle_seed, options=dataset_options) # pyrefly: ignore[bad-argument-type] process_split = _get_process_split( split, @@ -253,7 +253,7 @@ def get_data( process_count = jax.process_count() if rng_available: - rng = jax.random.fold_in(rng, process_index) # Derive RNG for this process. + rng = jax.random.fold_in(rng, process_index) # Derive RNG for this process. # pyrefly: ignore[bad-argument-type] rngs = list(jax.random.split(rng, 3)) else: rngs = 3 * [[None, None]] @@ -263,7 +263,7 @@ def get_data( data_dir=data_dir, split=split, shuffle_files=shuffle, - file_shuffle_seed=rngs.pop()[0], + file_shuffle_seed=rngs.pop()[0], # pyrefly: ignore[bad-argument-type] process_index=process_index, process_count=process_count, drop_remainder=drop_remainder) @@ -285,7 +285,7 @@ def get_data( if rng_available: ds = _preprocess_with_per_example_rng( - ds, preprocess_and_mask_fn, rng=rngs.pop()) + ds, preprocess_and_mask_fn, rng=rngs.pop()) # pyrefly: ignore[bad-argument-type] else: ds = ds.map(preprocess_and_mask_fn, num_parallel_calls=tf.data.AUTOTUNE) diff --git a/experimental/multimodal/preprocess_utils.py b/experimental/multimodal/preprocess_utils.py index a588bbcb5..419b366a2 100644 --- a/experimental/multimodal/preprocess_utils.py +++ b/experimental/multimodal/preprocess_utils.py @@ -373,7 +373,7 @@ def __call__(self, features: Features) -> Features: # than using tf.one_hot followed by tf.reduce_max; we tested. labels = features[self.key] if labels.shape.rank > 0 and self.multi: # pytype: disable=attribute-error # allow-recursive-types - x = tf.scatter_nd(labels[:, None], tf.ones(tf.shape(labels)[0]), + x = tf.scatter_nd(labels[:, None], tf.ones(tf.shape(labels)[0]), # pyrefly: ignore[bad-index] (self.depth,)) x = tf.clip_by_value(x, 0, 1) * (self.on - self.off) + self.off else: diff --git a/experimental/single_model_uncertainty/models/genomics_cnn.py b/experimental/single_model_uncertainty/models/genomics_cnn.py index 7a7760d78..a880c4fbe 100644 --- a/experimental/single_model_uncertainty/models/genomics_cnn.py +++ b/experimental/single_model_uncertainty/models/genomics_cnn.py @@ -56,7 +56,7 @@ def _conv_pooled_block( logging.info('conv input shape %s', inputs.shape) # expend last dim to fit conv2d out = tf.expand_dims(inputs, axis=3) # [batch_size, seq_len, embed_size, 1] - out = conv_layer(out) # [batch_size, seq_len, embed_size, num_motifs] + out = conv_layer(out) # [batch_size, seq_len, embed_size, num_motifs] # pyrefly: ignore[not-callable] # TODO(jjren) indices=1 only works for embed_size=4 out = tf.gather(out, indices=1, axis=2) # [batch_size, seq_len, num_motifs] @@ -73,7 +73,7 @@ def _input_embedding(inputs, vocab_size, one_hot=True, embed_size=None): embed_size = vocab_size else: if embed_size: - out = tf.keras.layers.Embedding(vocab_size, embed_size)(inputs) + out = tf.keras.layers.Embedding(vocab_size, embed_size)(inputs) # pyrefly: ignore[not-callable] else: raise ValueError('Embed input integers but embedding size is not given.') return out @@ -177,7 +177,7 @@ def genomics_cnn(batch_size: int, x, dropout_rate, use_mc_dropout, name='dropout2') if gp_layer_hparams and gp_layer_hparams['gp_input_dim'] > 0: # Uses random projection to reduce the input dimension of the GP layer. - x = tf.keras.layers.Dense( + x = tf.keras.layers.Dense( # pyrefly: ignore[not-callable] gp_layer_hparams['gp_input_dim'], kernel_initializer='random_normal', use_bias=False, @@ -185,4 +185,4 @@ def genomics_cnn(batch_size: int, name='gp_random_projection')( x) outputs = output_layer(num_classes, name='logits')(x) - return tf.keras.Model(inputs=inputs, outputs=outputs) + return tf.keras.Model(inputs=inputs, outputs=outputs) # pyrefly: ignore[bad-return] diff --git a/experimental/single_model_uncertainty/models/models.py b/experimental/single_model_uncertainty/models/models.py index 799d149cd..b4d5a9266 100644 --- a/experimental/single_model_uncertainty/models/models.py +++ b/experimental/single_model_uncertainty/models/models.py @@ -46,7 +46,7 @@ def get_model_names() -> List[str]: # edward2/tensorflow/initializers.py line 772. def get( model_name: str, - **hyperparameters) -> tf.keras.Model: + **hyperparameters) -> tf.keras.Model: # pyrefly: ignore[bad-return] """Gets a model builder by name. Args: diff --git a/experimental/single_model_uncertainty/models/util.py b/experimental/single_model_uncertainty/models/util.py index 5ee584449..c6dd5a773 100644 --- a/experimental/single_model_uncertainty/models/util.py +++ b/experimental/single_model_uncertainty/models/util.py @@ -173,9 +173,9 @@ def apply_dropout(inputs: tf.Tensor, dropout_rate, noise_shape=noise_shape, name=name) if use_mc_dropout: - return dropout_layer(inputs, training=True) + return dropout_layer(inputs, training=True) # pyrefly: ignore[not-callable] - return dropout_layer(inputs) + return dropout_layer(inputs) # pyrefly: ignore[not-callable] def make_output_layer(gp_layer_hparams=None): diff --git a/experimental/single_model_uncertainty/models/wide_resnet.py b/experimental/single_model_uncertainty/models/wide_resnet.py index 42fbc8676..e2e86178c 100644 --- a/experimental/single_model_uncertainty/models/wide_resnet.py +++ b/experimental/single_model_uncertainty/models/wide_resnet.py @@ -75,22 +75,22 @@ def basic_block( y = inputs y = BatchNormalization(beta_regularizer=tf.keras.regularizers.l2(l2), gamma_regularizer=tf.keras.regularizers.l2(l2))(y) - y = tf.keras.layers.Activation('relu')(y) + y = tf.keras.layers.Activation('relu')(y) # pyrefly: ignore[not-callable] y = models_util.apply_dropout(y, dropout_rate, use_mc_dropout) - y = conv_layer(filters=filters, + y = conv_layer(filters=filters, # pyrefly: ignore[not-callable] strides=strides, kernel_regularizer=tf.keras.regularizers.l2(l2))(y) y = BatchNormalization(beta_regularizer=tf.keras.regularizers.l2(l2), gamma_regularizer=tf.keras.regularizers.l2(l2))(y) - y = tf.keras.layers.Activation('relu')(y) + y = tf.keras.layers.Activation('relu')(y) # pyrefly: ignore[not-callable] y = models_util.apply_dropout(y, dropout_rate, use_mc_dropout) - y = conv_layer(filters=filters, + y = conv_layer(filters=filters, # pyrefly: ignore[not-callable] strides=1, kernel_regularizer=tf.keras.regularizers.l2(l2))(y) if not x.shape.is_compatible_with(y.shape): - x = conv_layer(filters=filters, + x = conv_layer(filters=filters, # pyrefly: ignore[not-callable] kernel_size=1, strides=strides, kernel_regularizer=tf.keras.regularizers.l2(l2))(x) @@ -209,15 +209,15 @@ def wide_resnet( conv_layer=conv2d) x = BatchNormalization(beta_regularizer=tf.keras.regularizers.l2(l2), gamma_regularizer=tf.keras.regularizers.l2(l2))(x) - x = tf.keras.layers.Activation('relu')(x) + x = tf.keras.layers.Activation('relu')(x) # pyrefly: ignore[not-callable] x = tf.keras.layers.AveragePooling2D(pool_size=8)(x) - x = tf.keras.layers.Flatten()(x) + x = tf.keras.layers.Flatten()(x) # pyrefly: ignore[not-callable] output_layer = models_util.make_output_layer( gp_layer_hparams=gp_layer_hparams) if gp_layer_hparams and gp_layer_hparams['gp_input_dim'] > 0: # Uses random projection to reduce the input dimension of the GP layer. - x = tf.keras.layers.Dense( + x = tf.keras.layers.Dense( # pyrefly: ignore[not-callable] gp_layer_hparams['gp_input_dim'], kernel_initializer='random_normal', use_bias=False,