Replace # pytype: disable suppressions with # pyrefly: ignore

enn/active_learning/priorities.py

@@ -40,7 +40,7 @@ def priority_fn(
         batch: datasets.ArrayBatch,
         key: chex.PRNGKey,
     ) -> base.PriorityOutput:
-      logits = enn_batch_fwd(params, state, batch.x)  # pytype: disable=wrong-arg-types  # numpy-scalars
+      logits = enn_batch_fwd(params, state, batch.x)  # pyrefly: ignore[bad-argument-type]
       # Make sure labels have shape [num_data, 1] as expected by priority.
       labels = batch.y
       if labels.ndim == 1:
@@ -49,7 +49,7 @@ def priority_fn(
       return values, {}
     return priority_fn
 
-  return make_priority_fn  # pytype: disable=bad-return-type  # numpy-scalars
+  return make_priority_fn  # pyrefly: ignore[bad-return]
 
 
 def uniform_per_example(

enn/datasets/utils.py

@@ -90,7 +90,7 @@ def _add_data_index(
     data_index: int, batch: ds_base.ArrayBatch
 ) -> ds_base.ArrayBatch:
   """Adds data_index into the batch."""
-  return ds_base.ArrayBatch(x=batch.x, y=batch.y, data_index=data_index)  # pytype: disable=wrong-arg-types  # numpy-scalars
+  return ds_base.ArrayBatch(x=batch.x, y=batch.y, data_index=data_index)  # pyrefly: ignore[bad-argument-type]
 
 
 class OverrideTrainDataset(ds_base.DatasetWithTransform):

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],  # pyrefly: ignore[invalid-type-var]
+      apply: base.ApplyFn[base.Input, base.Output],
       params: hk.Params,
       state: hk.State,
       batch: base.Data,
@@ -64,17 +64,17 @@ 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],  # pyrefly: ignore[invalid-type-var]
+  def loss_fn(enn: base.EpistemicNetwork[base.Input, base.Output],
               params: hk.Params,
               state: hk.State,
-              batch: base.Data,  # pyrefly: ignore[invalid-type-var]
+              batch: base.Data,
               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.
     batched_indexer = utils.make_batch_indexer(enn.indexer, num_index_samples)
     batched_loss = jax.vmap(single_loss, in_axes=[None, None, None, None, 0])
     loss, (new_state, metrics) = batched_loss(
-        enn.apply, params, state, batch, batched_indexer(key))  # pyrefly: ignore[bad-argument-type]
+        enn.apply, params, state, batch, batched_indexer(key))
 
     # Take the mean over the synthetic index batch dimension
     batch_mean = lambda x: jnp.mean(x, axis=0)
@@ -92,7 +92,7 @@ def loss_fn(enn: base.EpistemicNetwork[base.Input, base.Output],  # pyrefly: ign
 
     return mean_loss, (new_state, mean_metrics)
 
-  return loss_fn  # pyrefly: ignore[bad-return]
+  return loss_fn
 
 
 # Loss modules specialized to work only with Array inputs and Batch data.

enn/losses/prior_losses.py

@@ -177,7 +177,7 @@ def __call__(self, enn: networks.EnnArray,
       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
+    return loss, (state, {})
 
 
 @dataclasses.dataclass
@@ -212,4 +212,4 @@ def __call__(
       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
+    return loss, (state, {})

enn/losses/single_index.py

@@ -71,7 +71,7 @@ def __call__(
       batch: datasets.ArrayBatch,
       index: base.Index,
   ) -> base.LossOutput:
-    return self._loss(apply, params, state, batch, index)  # pyrefly: ignore[bad-argument-type]
+    return self._loss(apply, params, state, batch, index)
 
 
 def xent_loss_with_custom_labels(
@@ -103,7 +103,7 @@ def single_loss(
 
     loss = jnp.mean(batch_weights * softmax_xent)
     return loss, (state, {'loss': loss})
-  return single_loss  # pyrefly: ignore[bad-return]
+  return single_loss
 
 
 @dataclasses.dataclass
@@ -158,7 +158,7 @@ def __call__(
     chex.assert_equal_shape([log_likelihood, model_prior_kl])
     if self.temperature and self.input_dim:
       model_prior_kl *= jnp.sqrt(self.temperature) * self.input_dim
-    return model_prior_kl - log_likelihood, (state, {})  # pytype: disable=bad-return-type  # numpy-scalars
+    return model_prior_kl - log_likelihood, (state, {})  # pyrefly: ignore[bad-return]
 
 
 @dataclasses.dataclass
@@ -179,7 +179,7 @@ def __call__(
     net_out, state = apply(params, state, batch.x, index)
     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
+    return kl_term - log_likelihood, (state, {})  # pyrefly: ignore[bad-return]
 
 
 ################################################################################
@@ -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,  # pyrefly: ignore[bad-argument-type]
+    loss, (unused_state, metrics) = single_loss(apply_with_state, params,
                                                 constant_state, batch, index)
     return loss, metrics
 

enn/losses/utils.py

@@ -100,10 +100,10 @@ def combined_loss(
       apply: base.ApplyFn[base.Input, base.Output],
       params: hk.Params, state: hk.State, batch: base.Data,
       index: base.Index) -> base.LossOutput:
-    loss, (state, metrics) = train_loss(apply, params, state, batch, index)  # pyrefly: ignore[bad-argument-type]
+    loss, (state, metrics) = train_loss(apply, params, state, batch, index)
     for name, loss_fn in extra_losses.items():
       extra_loss, (unused_state,
-                   extra_metrics) = loss_fn(apply, params, state, batch, index)  # pyrefly: ignore[bad-argument-type]
+                   extra_metrics) = loss_fn(apply, params, state, batch, index)
       metrics[f'{name}:loss'] = extra_loss
       for key, value in extra_metrics.items():
         metrics[f'{name}:{key}'] = value

enn/losses/vae_losses.py

@@ -46,7 +46,7 @@ def binary_log_likelihood(x: chex.Array,
       x * output - jnp.logaddexp(0.0, output), axis=-1)
 
   chex.assert_shape(log_likelihood, (x.shape[0],))
-  return jnp.mean(log_likelihood)  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.mean(log_likelihood)  # pyrefly: ignore[bad-return]
 
 
 def gaussian_log_likelihood(x: chex.Array, mean: chex.Array,
@@ -75,7 +75,7 @@ def log_normal_prob(x: float, mu: float = 0, sigma: float = 1):
   log_likelihood = jnp.sum(log_likelihoods, axis=-1)
 
   chex.assert_shape(log_likelihood, (x.shape[0],))
-  return jnp.mean(log_likelihood)  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.mean(log_likelihood)  # pyrefly: ignore[bad-return]
 
 
 def latent_kl_divergence(mean: chex.Array,
@@ -95,7 +95,7 @@ def latent_kl_divergence(mean: chex.Array,
   kl = - 0.5 * jnp.sum(
       1. + log_var - jnp.square(mean) - jnp.exp(log_var), axis=-1)
   chex.assert_shape(kl, (mean.shape[0],))
-  return jnp.mean(kl)  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.mean(kl)  # pyrefly: ignore[bad-return]
 
 
 def latent_kl_fn(net_out: networks.OutputWithPrior) -> float:

enn/losses/vi_losses.py

@@ -109,7 +109,7 @@ def sum_log_scale_mixture_normal(
           loc=[mu_1, mu_2],  # One for each component.
           scale=[sigma_1, sigma_2]))
   log_probs = bimix_gauss.log_prob(x)
-  return jnp.sum(log_probs)  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.sum(log_probs)  # pyrefly: ignore[bad-return]
 
 
 def normal_log_prob(latent: chex.Array, sigma: float = 1, mu: float = 0):
@@ -154,7 +154,7 @@ def model_prior_kl_fn(out: networks.Output, params: hk.Params,
                         index: base.Index) -> float:
     """Compute the KL distance between model and prior densities using samples."""
     del index
-    latent = out.extra['hyper_index']  # pytype: disable=attribute-error  # numpy-scalars
+    latent = out.extra['hyper_index']  # pyrefly: ignore[missing-attribute]
 
     # Calculate prior
     log_priors_sum = sum_log_scale_mixture_normal(latent, sigma_1,
@@ -177,7 +177,7 @@ def model_prior_kl_fn(out: networks.Output, params: hk.Params,
     log_var_posteriors = log_normal_prob_vectorized(latent, mus, sigmas)
     log_var_posteriors_sum = jnp.sum(log_var_posteriors)
 
-    return (log_var_posteriors_sum - log_priors_sum) / num_samples  # pytype: disable=bad-return-type  # jnp-type
+    return (log_var_posteriors_sum - log_priors_sum) / num_samples  # pyrefly: ignore[bad-return]
 
   return model_prior_kl_fn
 
@@ -226,7 +226,7 @@ def model_prior_kl_fn(out: networks.Output, params: hk.Params,
     weights = jnp.concatenate(weights, axis=0)
     scales = jnp.log(1 + jnp.exp(weights))
     chex.assert_equal_shape([scales, biases])
-    return 0.5  / num_samples * (  # pytype: disable=bad-return-type  # jnp-type
+    return 0.5  / num_samples * (  # pyrefly: ignore[bad-return]
         jnp.sum(jnp.square(scales)) / (sigma_0 ** 2)
         + jnp.sum(jnp.square(biases)) / (sigma_0 ** 2)
         - len(biases)
@@ -292,7 +292,7 @@ def model_prior_kl_fn(out: networks.Output, params: hk.Params,
     weights_biases = weights @ biases
     chex.assert_equal(len(weights_biases), index_dim)
     proj_biases_norm = weights_biases @ w_sq_inv @ weights_biases.T
-    return 0.5  / num_samples * (sigma_u_trace - index_dim - sigma_u_log_det  # pytype: disable=bad-return-type  # jnp-type
+    return 0.5  / num_samples * (sigma_u_trace - index_dim - sigma_u_log_det
                                  + proj_biases_norm / sigma_0**2)
 
   return model_prior_kl_fn

enn/metrics/calibration.py

@@ -266,4 +266,4 @@ def _map_stats_to_ece(ece_stats: Dict[str, chex.Array]) -> float:
   accs = jnp.nan_to_num(ece_stats['correct_sums'] / counts)
   confs = jnp.nan_to_num(ece_stats['prob_sums'] / counts)
   total_count = jnp.sum(counts)
-  return jnp.sum(counts/ total_count * jnp.abs(accs - confs))  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.sum(counts/ total_count * jnp.abs(accs - confs))  # pyrefly: ignore[bad-return]

enn/metrics/joint.py

@@ -45,7 +45,7 @@ def joint_ll_repeat(logits: chex.Array,
 
     # Sample with replacement from the anchor points and sum for joint ll
     selected = jax.random.randint(key, shape=[tau], minval=0, maxval=kappa)
-    return jnp.sum(log_probs[selected])  # pytype: disable=bad-return-type  # jnp-type
+    return jnp.sum(log_probs[selected])  # pyrefly: ignore[bad-return]
 
   def enn_nll(logits: chex.Array,
               labels: chex.Array,
@@ -58,7 +58,7 @@ def enn_nll(logits: chex.Array,
     # Averaging over ENN samples
     batched_ll = jax.vmap(joint_ll_repeat, in_axes=[0, None, None])
     lls = batched_ll(logits, labels, key)  # pyrefly: ignore[bad-argument-type]
-    return -1 * metrics_base.average_sampled_log_likelihood(lls)  # pytype: disable=wrong-arg-types  # numpy-scalars
+    return -1 * metrics_base.average_sampled_log_likelihood(lls)  # pyrefly: ignore[bad-argument-type]
 
   def polyadic_nll(logits: chex.Array, labels: chex.Array) -> float:
     """Returns polyadic NLL based on repeated inputs.
@@ -90,7 +90,7 @@ def polyadic_nll(logits: chex.Array, labels: chex.Array) -> float:
         logits, labels, batch_size=kappa)
     keys = jax.random.split(jax.random.PRNGKey(seed), batched_logits.shape[0])
     nlls = jax.vmap(enn_nll, in_axes=0)(batched_logits, batched_labels, keys)
-    return jnp.mean(nlls)  # pytype: disable=bad-return-type  # jnp-type
+    return jnp.mean(nlls)  # pyrefly: ignore[bad-return]
 
   return jax.jit(polyadic_nll)
 
@@ -110,7 +110,7 @@ def calculate_nll_joint(logits: chex.Array, labels: chex.Array) -> float:
     lls = jax.vmap(calculate_joint_ll)(
         batched_logits, batched_labels)
     chex.assert_shape(lls, (num_batches,))
-    return -1 * jnp.mean(lls)  # pytype: disable=bad-return-type  # jnp-type
+    return -1 * jnp.mean(lls)  # pyrefly: ignore[bad-return]
 
   return calculate_nll_joint
 
@@ -137,7 +137,7 @@ def calculate_joint_ll(logits: chex.Array, labels: chex.Array) -> float:
 
   batched_ll = jax.vmap(marginal.categorical_log_likelihood, in_axes=[0, None])
   sampled_ll = batched_ll(class_probs, labels)
-  return metrics_base.average_sampled_log_likelihood(sampled_ll)  # pytype: disable=wrong-arg-types  # numpy-scalars
+  return metrics_base.average_sampled_log_likelihood(sampled_ll)  # pyrefly: ignore[bad-argument-type]
 
 
 def reshape_to_smaller_batches(

enn/metrics/marginal.py

@@ -56,12 +56,12 @@ def calculate_accuracy(logits: chex.Array, labels: chex.Array) -> float:
   predictions = jnp.argmax(mean_class_prob, axis=1)[:, None]
   chex.assert_shape(predictions, [num_data, 1])
 
-  return jnp.mean(predictions == labels)  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.mean(predictions == labels)  # pyrefly: ignore[bad-return]
 
 
 def categorical_log_likelihood(probs: chex.Array, labels: chex.Array) -> float:
   """Computes joint log likelihood based on probs and labels."""
   num_data, unused_num_classes = probs.shape
   assert len(labels) == num_data  # pyrefly: ignore[bad-argument-type]
   assigned_probs = probs[jnp.arange(num_data), jnp.squeeze(labels)]  # pyrefly: ignore[bad-index]
-  return jnp.sum(jnp.log(assigned_probs))  # pytype: disable=bad-return-type  # jnp-type
+  return jnp.sum(jnp.log(assigned_probs))  # pyrefly: ignore[bad-return]

enn/networks/bert/bert.py

@@ -56,7 +56,7 @@ def net_fn(inputs: base.BertInput) -> networks_base.OutputWithPrior:
     )
 
     # Embed and summarize the sequence.
-    return bert_model(  # pytype: disable=wrong-arg-types  # jax-devicearray
+    return bert_model(
         input_ids=inputs.token_ids,
         token_type_ids=inputs.segment_ids,
         input_mask=inputs.input_mask.astype(jnp.int32),

enn/networks/categorical_ensembles.py

@@ -52,7 +52,7 @@ def __call__(self, inputs: chex.Array) -> chex.Array:
     """Apply MLP and wrap outputs appropriately."""
     out = hk.Flatten()(inputs)
     out = hk.nets.MLP(self.output_sizes)(out)
-    return CatOutputWithPrior(  # pytype: disable=bad-return-type  # numpy-scalars
+    return CatOutputWithPrior(  # pyrefly: ignore[bad-return]
         train=jnp.reshape(out, [-1, self.dim_out, len(self.atoms)]),
         extra={'atoms': self.atoms},
     )

enn/networks/combiners.py

@@ -145,4 +145,4 @@ def sample_logits(sub_key: chex.PRNGKey) -> chex.Array:
 
     return jax.vmap(sample_logits)(enn_keys)
 
-  return jax.jit(enn_batch_fwd)  # pyrefly: ignore[bad-return]
+  return jax.jit(enn_batch_fwd)

enn/networks/einsum_mlp.py

@@ -51,7 +51,7 @@ def make_einsum_ensemble_mlp_enn(
   def ensemble_forward(x: chex.Array) -> networks_base.OutputWithPrior:
     """Forwards the entire ensemble at given input x."""
     model = EnsembleMLP(output_sizes, num_ensemble, nonzero_bias, activation)
-    return model(x)  # pytype: disable=bad-return-type  # jax-ndarray
+    return model(x)
 
   transformed = hk.without_apply_rng(hk.transform(ensemble_forward))
 
@@ -60,7 +60,7 @@ def apply(params: hk.Params, x: chex.Array,
             z: base.Index) -> networks_base.OutputWithPrior:
     net_out = transformed.apply(params, x)
     one_hot_index = jax.nn.one_hot(z, num_ensemble)
-    return jnp.dot(net_out, one_hot_index)  # pytype: disable=bad-return-type  # jnp-type
+    return jnp.dot(net_out, one_hot_index)  # pyrefly: ignore[bad-return]
 
   def init(key: chex.PRNGKey, x: chex.Array,
            z: base.Index) -> hk.Params:
@@ -237,4 +237,4 @@ def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:  # [B, H] -> [B, D, K]
       out = layer(out)
       if i < num_layers - 1:
         out = self.activation(out)
-    return out  # pytype: disable=bad-return-type  # numpy-scalars
+    return out

enn/networks/forwarders.py

@@ -58,4 +58,4 @@ def forward(params: hk.Params, state: hk.State, x: base.Input) -> chex.Array:
     net_out, unused_state = batch_apply(params, state, x, indices)
     return utils.parse_net_output(net_out)
 
-  return jax.jit(forward)  # pyrefly: ignore[bad-return]
+  return jax.jit(forward)

enn/networks/gaussian_enn.py

@@ -79,7 +79,7 @@ def __init__(self,
 
     def net_fn(inputs: chex.Array) -> chex.Array:
       with hke.custom_getter(enn_getter), hke.custom_creator(enn_creator):
-        output = module_ctor()(inputs)  # pytype: disable=not-callable
+        output = module_ctor()(inputs)  # pyrefly: ignore[not-callable]
         return output
 
     # TODO(author2): Note that the GaussianENN requires a rng_key in place of an

enn/networks/hypermodels.py

@@ -68,7 +68,6 @@ def hyper_torso(index):
     super().__init__(enn.apply, enn.init, enn.indexer)
 
 
-# pytype: disable=bad-return-type
 def hypermodel_module(
     transformed_base: hk.Transformed,
     dummy_input: chex.Array,
@@ -177,14 +176,13 @@ def hyper_fn(inputs: chex.Array,
           prior=jnp.zeros_like(out),
           extra={
               'hyper_net_out': generated_params,
-              'base_net_params': generated_params_scaled,
+              'base_net_params': generated_params_scaled,  # pyrefly: ignore[bad-assignment]
               'hyper_index': hyper_index
           })
     return out
 
   enn_module = hk.to_module(hyper_fn)
-  return enn_module
-# pytype: enable=bad-return-type
+  return enn_module  # pyrefly: ignore[bad-return]
 
 
 class MLPHypermodelWithHypermodelPrior(networks_base.EnnArray):

enn/networks/priors.py

@@ -207,6 +207,6 @@ class NetworkWithAdditivePrior(hk.Module):
 
   def __call__(self, *args, **kwargs) -> chex.Array:
     logging.warning(WARN)
-    prior = jax.lax.stop_gradient(self.prior_net(*args, **kwargs))  # pytype:disable=not-callable
-    net_out = self.net(*args, **kwargs)  # pytype:disable=not-callable
+    prior = jax.lax.stop_gradient(self.prior_net(*args, **kwargs))  # pyrefly: ignore[not-callable]
+    net_out = self.net(*args, **kwargs)  # pyrefly: ignore[not-callable]
     return net_out + prior * self.prior_scale