Automated Code Change

enn/active_learning/priorities.py

@@ -230,7 +230,7 @@ def compute_ucb(
     chex.assert_shape(ucb_value, (data_size,))
     return ucb_value
 
-  return compute_ucb
+  return compute_ucb  # pyrefly: ignore[bad-return]
 
 
 def make_scaled_mean_per_example(
@@ -257,7 +257,7 @@ def compute_scaled_mean(
     chex.assert_shape(mean_values, (data_size,))
     return mean_values
 
-  return compute_scaled_mean
+  return compute_scaled_mean  # pyrefly: ignore[bad-return]
 
 
 def make_scaled_std_per_example(
@@ -284,7 +284,7 @@ def compute_scaled_std(
     chex.assert_shape(std_values, (data_size,))
     return std_values
 
-  return compute_scaled_std
+  return compute_scaled_std  # pyrefly: ignore[bad-return]
 
 
 _PerExamplePriorities = {
@@ -303,7 +303,7 @@ def compute_scaled_std(
 
 
 _PriorityFnCtors = {
-    key: make_priority_fn_ctor(value)
+    key: make_priority_fn_ctor(value)  # pyrefly: ignore[bad-argument-type]
     for key, value in _PerExamplePriorities.items()
 }
 
@@ -327,4 +327,4 @@ def get_implemented_per_example_priorities() -> tp.Sequence[str]:
 def get_per_example_priority(name: str) -> base.PerExamplePriority:
   """Returns a per example priority function for the priority specified by `name`."""
   assert name in get_implemented_per_example_priorities()
-  return _PerExamplePriorities[name]
+  return _PerExamplePriorities[name]  # pyrefly: ignore[bad-return]

enn/active_learning/prioritized.py

@@ -51,7 +51,7 @@ def sample_batch(
       key: chex.PRNGKey,
   ) -> datasets.ArrayBatch:
     """Acquires data. This is the function per device (can get pmaped)."""
-    pool_size = len(batch.y)
+    pool_size = len(batch.y)  # pyrefly: ignore[bad-argument-type]
 
     candidate_scores, unused_metrics = self._priority_fn(
         params, state, batch, key)

enn/data_noise/bootstrapping.py

@@ -48,7 +48,7 @@ def __call__(self, data: datasets.ArrayBatch,
                index: base.Index) -> datasets.ArrayBatch:
     """Apply bootstrap reweighting to a batch of data."""
     boot_fn = make_boot_fn(self.enn, self.distribution, self.seed)
-    boot_weights = boot_fn(data.data_index, index)
+    boot_weights = boot_fn(data.data_index, index)  # pyrefly: ignore[bad-argument-type]
     return dataclasses.replace(data, weights=boot_weights)
 
 
@@ -110,7 +110,7 @@ def make_boot_fn(enn: _ENN,
     if distribution not in DISTRIBUTIONS:
       raise ValueError(f'dist={distribution} not implemented for ensemble.')
     weight_fn = DISTRIBUTIONS[distribution]
-    return _make_ensemble_bootstrap_fn(weight_fn, seed)
+    return _make_ensemble_bootstrap_fn(weight_fn, seed)  # pyrefly: ignore[bad-argument-type]
 
   # Bootstrapping for Gaussian with unit index
   elif isinstance(indexer, networks.GaussianWithUnitIndexer):
@@ -130,7 +130,7 @@ def make_boot_fn(enn: _ENN,
       return _make_gaussian_index_bernoulli_bootstrap(index_dim, seed)
     elif distribution == 'exponential':
       return _make_gaussian_index_exponential_bootstrap(
-          index_dim, seed, 1/jnp.sqrt(index_dim))
+          index_dim, seed, 1/jnp.sqrt(index_dim))  # pyrefly: ignore[bad-argument-type]
     else:
       raise ValueError(
           f'dist={distribution} not implemented for GaussianIndexer.')
@@ -151,7 +151,7 @@ def make_boot_fn(enn: _ENN,
     if distribution not in DISTRIBUTIONS:
       raise ValueError(f'dist={distribution} not implemented for gauss_enn.')
     weight_fn = DISTRIBUTIONS[distribution]
-    return _make_prng_bootstrap_fn(weight_fn)
+    return _make_prng_bootstrap_fn(weight_fn)  # pyrefly: ignore[bad-argument-type]
 
   else:
     raise ValueError(
@@ -193,7 +193,7 @@ def boot_fn(data_index: datasets.DataIndex, index: base.Index):
     """Assumes integer index for ensemble weights."""
     chex.assert_shape(data_index, (None, 1))
     if not index.shape:  # If it's a single integer -> repeat for batch
-      index = jnp.repeat(index, len(data_index))
+      index = jnp.repeat(index, len(data_index))  # pyrefly: ignore[bad-argument-type]
     data_keys = _make_key(data_index, seed)
     rng_keys = fold_in(data_keys, index)
     return weight_fn(rng_keys)[:, None]

enn/data_noise/gaussian.py

@@ -43,7 +43,7 @@ def __call__(self, data: datasets.ArrayBatch,
     """Apply Gaussian noise to the target y."""
     chex.assert_shape(data.y, (None, 1))  # Only implemented for 1D now.
     noise_fn = make_noise_fn(self.enn, self.noise_std, self.seed)
-    y_noise = noise_fn(data.data_index, index)
+    y_noise = noise_fn(data.data_index, index)  # pyrefly: ignore[bad-argument-type]
     return dataclasses.replace(data, y=data.y + y_noise)
 
 
@@ -94,7 +94,7 @@ def noise_fn(data_index: datasets.DataIndex,
     """Assumes integer index for ensemble."""
     chex.assert_shape(data_index, (None, 1))
     if not index.shape:  # If it's a single integer -> repeat for batch
-      index = jnp.repeat(index, len(data_index))
+      index = jnp.repeat(index, len(data_index))  # pyrefly: ignore[bad-argument-type]
     data_keys = _make_key(data_index, seed)
     batch_keys = batch_fold_in(data_keys, index)
     samples = batch_normal(batch_keys)[:, None]

enn/datasets/base.py

@@ -88,9 +88,9 @@ class DatasetWithTransform(Dataset):
 
 
 class OodVariant(enum.Enum):
-  WHOLE: str = 'eval'
-  IN_DISTRIBUTION: str = 'eval_in_dist'
-  OUT_DISTRIBUTION: str = 'eval_out_dist'
+  WHOLE: str = 'eval'  # pyrefly: ignore[invalid-annotation]
+  IN_DISTRIBUTION: str = 'eval_in_dist'  # pyrefly: ignore[invalid-annotation]
+  OUT_DISTRIBUTION: str = 'eval_out_dist'  # pyrefly: ignore[invalid-annotation]
 
   @classmethod
   def valid_values(cls) -> tp.List[str]:

enn/datasets/cifar.py

@@ -69,7 +69,7 @@ class Cifar(ds_base.DatasetWithTransform):
   num_train: int = 50_000
   train_ds_transformer: ds_base.DatasetTransformer = lambda x: x
   eval_ds_transformers: Dict[
-      str, ds_base.DatasetTransformer] = ds_base.EVAL_TRANSFORMERS_DEFAULT
+      str, ds_base.DatasetTransformer] = ds_base.EVAL_TRANSFORMERS_DEFAULT  # pyrefly: ignore[bad-assignment]
   # Whether to add a leading axis of number of devices to the batches. If true,
   # data batches have shape (number_devices, batch_size / number_devices, ...);
   # otherwise, they have shape of (batch_size, ...).

enn/datasets/imagenet.py

@@ -59,7 +59,7 @@ class Imagenet(ds_base.DatasetWithTransform):
   num_train: int = Split.TRAIN_AND_VALID.num_examples
   train_ds_transformer: ds_base.DatasetTransformer = lambda x: x
   eval_ds_transformers: Dict[
-      str, ds_base.DatasetTransformer] = ds_base.EVAL_TRANSFORMERS_DEFAULT
+      str, ds_base.DatasetTransformer] = ds_base.EVAL_TRANSFORMERS_DEFAULT  # pyrefly: ignore[bad-assignment]
   # Whether to add a leading axis of number of devices to the batches. If true,
   # data batches have shape (number_devices, batch_size / number_devices, ...);
   # otherwise, they have shape of (batch_size, ...).
@@ -171,7 +171,7 @@ def load(
     yield from it.repeat(batch, end - start)
     return
 
-  total_batch_size = np.prod(batch_dims)
+  total_batch_size = np.prod(batch_dims)  # pyrefly: ignore[bad-assignment]
 
   tfds_split = tfds.core.ReadInstruction(
       _to_tfds_split(split), from_=start, to=end, unit='abs')
@@ -200,7 +200,7 @@ def load(
       # Only cache if we are reading a subset of the dataset.
       ds = ds.cache()
     ds = ds.repeat()
-    seed_shuffle = tf.cast(rngs[0][0], tf.int64) if seed is not None else 0
+    seed_shuffle = tf.cast(rngs[0][0], tf.int64) if seed is not None else 0  # pyrefly: ignore[unbound-name]
     ds = ds.shuffle(buffer_size=10 * total_batch_size, seed=seed_shuffle)
 
   else:
@@ -276,10 +276,10 @@ def _to_tfds_split(split: Split) -> tfds.Split:
   # competition, so it has been typical at DeepMind to consider the VALID
   # split the TEST split and to reserve 10k images from TRAIN for VALID.
   if split in (Split.TRAIN, Split.TRAIN_AND_VALID, Split.VALID):
-    return tfds.Split.TRAIN
+    return tfds.Split.TRAIN  # pyrefly: ignore[missing-attribute]
   else:
     assert split == Split.TEST
-    return tfds.Split.VALIDATION
+    return tfds.Split.VALIDATION  # pyrefly: ignore[missing-attribute]
 
 
 def _shard(split: Split,

enn/datasets/mnist.py

@@ -48,7 +48,7 @@ class Mnist(ds_base.DatasetWithTransform):
   normalization_mode: str = 'standard'
   train_ds_transformer: ds_base.DatasetTransformer = lambda x: x
   eval_ds_transformers: Dict[
-      str, ds_base.DatasetTransformer] = ds_base.EVAL_TRANSFORMERS_DEFAULT
+      str, ds_base.DatasetTransformer] = ds_base.EVAL_TRANSFORMERS_DEFAULT  # pyrefly: ignore[bad-assignment]
   # Whether to add a leading axis of number of devices to the batches. If true,
   # data batches have shape (number_devices, batch_size / number_devices, ...);
   # otherwise, they have shape of (batch_size, ...).

enn/extra/kmeans.py

@@ -43,7 +43,7 @@ def fit(self, x: chex.Array) -> KMeansOutput:
     # Initialize centroids randomly
     random_idx = jax.random.choice(
         self.key, x.shape[0], [self.num_centroids], replace=False)
-    initial_centroids = x[random_idx, :]
+    initial_centroids = x[random_idx, :]  # pyrefly: ignore[bad-index]
     initial_state = _TrainingState(initial_centroids, iter=0)
 
     # Perfom KMeans via jax.lax.while_loop
@@ -91,7 +91,7 @@ def kmeans_iteration(x: chex.Array, state: _TrainingState) -> _TrainingState:
   chex.assert_shape(one_hot_centroids, [num_x, num_centroids])
 
   # Take mean over classes for new centroids.
-  masked_x = x[:, None, :] * one_hot_centroids[:, :, None]
+  masked_x = x[:, None, :] * one_hot_centroids[:, :, None]  # pyrefly: ignore[bad-index]
   chex.assert_shape(masked_x, [num_x, num_centroids, dim_x])
   sum_per_centroid = jnp.sum(masked_x, axis=0)
   count_per_centroid = jnp.sum(one_hot_centroids, axis=0)

enn/losses/base.py

@@ -37,7 +37,7 @@ class SingleLossFn(te.Protocol[base.Input, base.Output, base.Data,]):
 
   def __call__(
       self,
-      apply: base.ApplyFn[base.Input, base.Output],
+      apply: base.ApplyFn[base.Input, base.Output],  # pyrefly: ignore[invalid-type-var]
       params: hk.Params,
       state: hk.State,
       batch: base.Data,
@@ -64,10 +64,10 @@ def average_single_index_loss(
     LossFn that comprises the mean of both the loss and the metrics.
   """
 
-  def loss_fn(enn: base.EpistemicNetwork[base.Input, base.Output],
+  def loss_fn(enn: base.EpistemicNetwork[base.Input, base.Output],  # pyrefly: ignore[invalid-type-var]
               params: hk.Params,
               state: hk.State,
-              batch: base.Data,
+              batch: base.Data,  # pyrefly: ignore[invalid-type-var]
               key: chex.PRNGKey) -> base.LossOutput:
     # Apply the loss in parallel over num_index_samples different indices.
     # This is the key logic to this loss function.
@@ -92,7 +92,7 @@ def loss_fn(enn: base.EpistemicNetwork[base.Input, base.Output],
 
     return mean_loss, (new_state, mean_metrics)
 
-  return loss_fn
+  return loss_fn  # pyrefly: ignore[bad-return]
 
 
 # Loss modules specialized to work only with Array inputs and Batch data.

enn/losses/categorical_regression.py

@@ -34,10 +34,10 @@ def transform_to_2hot(target: chex.Array,
   """Converts a scalar target to a 2-hot encoding of the nearest support."""
   target = jnp.clip(target, support.min(), support.max())
   high_idx = jnp.sum(support < target)
-  num_bins = len(support)
+  num_bins = len(support)  # pyrefly: ignore[bad-argument-type]
 
-  low_value = support[high_idx - 1]
-  high_value = support[high_idx]
+  low_value = support[high_idx - 1]  # pyrefly: ignore[bad-index]
+  high_value = support[high_idx]  # pyrefly: ignore[bad-index]
   prob = (target - high_value) / (low_value - high_value)
 
   lower_one_hot = prob * rlax.one_hot(high_idx - 1, num_bins)
@@ -63,15 +63,15 @@ def __call__(self, apply: networks.ApplyArray,
     assert isinstance(net_out, networks.CatOutputWithPrior)
 
     # Form the target values in real space
-    target_val = batch.y - net_out.prior
+    target_val = batch.y - net_out.prior  # pyrefly: ignore[unsupported-operation]
 
     # Convert values to 2-hot target probabilities
     probs = jax.vmap(transform_to_2hot, in_axes=[0, None])(
         jnp.squeeze(target_val), net_out.extra['atoms'])
     probs = jnp.expand_dims(probs, 1)
     xent_loss = -jnp.sum(probs * jax.nn.log_softmax(net_out.train), axis=-1)
     if batch.weights is None:
-      batch_weights = jnp.ones_like(batch.data_index)
+      batch_weights = jnp.ones_like(batch.data_index)  # pyrefly: ignore[bad-argument-type]
     else:
       batch_weights = batch.weights
     chex.assert_equal_shape([batch_weights, xent_loss])

enn/losses/prior_losses.py

@@ -174,9 +174,9 @@ def __call__(self, enn: networks.EnnArray,
 
     # Distill aggregate stats to the "mean_index"
     if hasattr(enn.indexer, 'mean_index') and self.distill_index:
-      distill_out = enn.apply(params, fake_x, enn.indexer.mean_index)
-      loss += distill_mean_regression(batched_out, distill_out)
-      loss += distill_var_regression(batched_out, distill_out)
+      distill_out = enn.apply(params, fake_x, enn.indexer.mean_index)  # pyrefly: ignore[bad-argument-type, missing-argument]
+      loss += distill_mean_regression(batched_out, distill_out)  # pyrefly: ignore[bad-argument-type]
+      loss += distill_var_regression(batched_out, distill_out)  # pyrefly: ignore[bad-argument-type]
     return loss, (state, {})  # pytype: disable=bad-return-type  # numpy-scalars
 
 
@@ -209,7 +209,7 @@ def __call__(
 
     # Distill aggregate stats to the "mean_index"
     if hasattr(enn.indexer, 'mean_index') and self.distill_index:
-      distill_out = enn.apply(params, fake_x, enn.indexer.mean_index)
-      loss += distill_mean_classification(batched_out, distill_out)
-      loss += distill_var_classification(batched_out, distill_out)
+      distill_out = enn.apply(params, fake_x, enn.indexer.mean_index)  # pyrefly: ignore[bad-argument-type, missing-argument]
+      loss += distill_mean_classification(batched_out, distill_out)  # pyrefly: ignore[bad-argument-type]
+      loss += distill_var_classification(batched_out, distill_out)  # pyrefly: ignore[bad-argument-type]
     return loss, (state, {})  # pytype: disable=bad-return-type  # numpy-scalars

enn/losses/single_index.py

@@ -44,9 +44,9 @@ def __call__(self,
     net_out, state = apply(params, state, batch.x, index)
     net_out = networks.parse_net_output(net_out)
     chex.assert_equal_shape([net_out, batch.y])
-    sq_loss = jnp.square(networks.parse_net_output(net_out) - batch.y)
+    sq_loss = jnp.square(networks.parse_net_output(net_out) - batch.y)  # pyrefly: ignore[unsupported-operation]
     if batch.weights is None:
-      batch_weights = jnp.ones_like(batch.data_index)
+      batch_weights = jnp.ones_like(batch.data_index)  # pyrefly: ignore[bad-argument-type]
     else:
       batch_weights = batch.weights
     chex.assert_equal_shape([batch_weights, sq_loss])
@@ -71,7 +71,7 @@ def __call__(
       batch: datasets.ArrayBatch,
       index: base.Index,
   ) -> base.LossOutput:
-    return self._loss(apply, params, state, batch, index)
+    return self._loss(apply, params, state, batch, index)  # pyrefly: ignore[bad-argument-type]
 
 
 def xent_loss_with_custom_labels(
@@ -90,7 +90,7 @@ def single_loss(
     chex.assert_shape(batch.y, (None, 1))
     net_out, state = apply(params, state, batch.x, index)
     logits = networks.parse_net_output(net_out)
-    labels = labeller(batch.y[:, 0])
+    labels = labeller(batch.y[:, 0])  # pyrefly: ignore[bad-index]
 
     softmax_xent = -jnp.sum(
         labels * jax.nn.log_softmax(logits), axis=1, keepdims=True)
@@ -103,7 +103,7 @@ def single_loss(
 
     loss = jnp.mean(batch_weights * softmax_xent)
     return loss, (state, {'loss': loss})
-  return single_loss
+  return single_loss  # pyrefly: ignore[bad-return]
 
 
 @dataclasses.dataclass
@@ -120,7 +120,7 @@ def __call__(self, apply: networks.ApplyArray,
     net_out, state = apply(params, state, batch.x, index)
     logits = networks.parse_net_output(net_out)
     preds = jnp.argmax(logits, axis=1)
-    correct = (preds == batch.y[:, 0])
+    correct = (preds == batch.y[:, 0])  # pyrefly: ignore[bad-index]
     accuracy = jnp.mean(correct)
     return 1 - accuracy, (state, {'accuracy': accuracy})
 
@@ -177,8 +177,8 @@ def __call__(
       index: base.Index,
   ) -> base.LossOutput:
     net_out, state = apply(params, state, batch.x, index)
-    kl_term = self.latent_kl_fn(net_out)
-    log_likelihood = self.log_likelihood_fn(net_out, batch)
+    kl_term = self.latent_kl_fn(net_out)  # pyrefly: ignore[bad-argument-type]
+    log_likelihood = self.log_likelihood_fn(net_out, batch)  # pyrefly: ignore[bad-argument-type]
     return kl_term - log_likelihood, (state, {})  # pytype: disable=bad-return-type  # numpy-scalars
 
 
@@ -202,7 +202,7 @@ def new_loss(
       index: base.Index,
   ) -> losses_base.LossOutputNoState:
     apply_with_state = networks.wrap_apply_no_state_as_apply(apply)
-    loss, (unused_state, metrics) = single_loss(apply_with_state, params,
+    loss, (unused_state, metrics) = single_loss(apply_with_state, params,  # pyrefly: ignore[bad-argument-type]
                                                 constant_state, batch, index)
     return loss, metrics