diff --git a/enn/active_learning/priorities.py b/enn/active_learning/priorities.py
index 49ea0b9..422efb0 100644
--- a/enn/active_learning/priorities.py
+++ b/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]
diff --git a/enn/active_learning/prioritized.py b/enn/active_learning/prioritized.py
index 7380abc..898d792 100644
--- a/enn/active_learning/prioritized.py
+++ b/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)
diff --git a/enn/data_noise/bootstrapping.py b/enn/data_noise/bootstrapping.py
index 4ae895c..6399369 100644
--- a/enn/data_noise/bootstrapping.py
+++ b/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]
diff --git a/enn/data_noise/gaussian.py b/enn/data_noise/gaussian.py
index 676c974..ff96859 100644
--- a/enn/data_noise/gaussian.py
+++ b/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]
diff --git a/enn/datasets/base.py b/enn/datasets/base.py
index 3b62f7d..bb5a6bf 100644
--- a/enn/datasets/base.py
+++ b/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]:
diff --git a/enn/datasets/cifar.py b/enn/datasets/cifar.py
index 8993050..3270a3f 100644
--- a/enn/datasets/cifar.py
+++ b/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, ...).
diff --git a/enn/datasets/imagenet.py b/enn/datasets/imagenet.py
index b800fff..a8608c0 100644
--- a/enn/datasets/imagenet.py
+++ b/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,
diff --git a/enn/datasets/mnist.py b/enn/datasets/mnist.py
index 4a5f4b8..56b9f47 100644
--- a/enn/datasets/mnist.py
+++ b/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, ...).
diff --git a/enn/extra/kmeans.py b/enn/extra/kmeans.py
index 20db8a7..ab3fdf1 100644
--- a/enn/extra/kmeans.py
+++ b/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)
diff --git a/enn/losses/base.py b/enn/losses/base.py
index 894e18a..74043a6 100644
--- a/enn/losses/base.py
+++ b/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.
diff --git a/enn/losses/categorical_regression.py b/enn/losses/categorical_regression.py
index e188a2c..13ace33 100644
--- a/enn/losses/categorical_regression.py
+++ b/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,7 +63,7 @@ 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])(
@@ -71,7 +71,7 @@ def __call__(self, apply: networks.ApplyArray,
     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])
diff --git a/enn/losses/prior_losses.py b/enn/losses/prior_losses.py
index a936d73..5b34985 100644
--- a/enn/losses/prior_losses.py
+++ b/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
diff --git a/enn/losses/single_index.py b/enn/losses/single_index.py
index 9ea43a5..2ded1a8 100644
--- a/enn/losses/single_index.py
+++ b/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
 
diff --git a/enn/losses/utils.py b/enn/losses/utils.py
index cf87c22..f898f46 100644
--- a/enn/losses/utils.py
+++ b/enn/losses/utils.py
@@ -44,7 +44,7 @@ def l2_weights_with_predicate(
   """Sum of squares of parameter weights that passes predicate_fn."""
   if predicate is not None:
     params = hk.data_structures.filter(predicate, params)
-  return sum(jnp.sum(jnp.square(p)) for p in jax.tree.leaves(params))
+  return sum(jnp.sum(jnp.square(p)) for p in jax.tree.leaves(params))  # pyrefly: ignore[bad-return]
 
 
 def add_data_noise(
@@ -54,7 +54,7 @@ def add_data_noise(
   """Applies a DataNoise function to each batch of data."""
 
   def noisy_loss(
-      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,
@@ -72,7 +72,7 @@ def add_l2_weight_decay(
 ) -> _LossFn:
   """Adds scale * l2 weight decay to an existing loss function."""
   try:  # Scale is numeric.
-    scale = jnp.sqrt(scale)
+    scale = jnp.sqrt(scale)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     scale_fn = lambda ps: jax.tree_util.tree_map(lambda p: scale * p, ps)
   except TypeError:
     scale_fn = scale  # Assuming scale is a Callable.
@@ -82,7 +82,7 @@ def new_loss(
       params: hk.Params, state: hk.State, batch: base.Data,
       key: chex.PRNGKey) -> base.LossOutput:
     loss, (state, metrics) = loss_fn(enn, params, state, batch, key)
-    decay = l2_weights_with_predicate(scale_fn(params), predicate)
+    decay = l2_weights_with_predicate(scale_fn(params), predicate)  # pyrefly: ignore[not-callable]
     total_loss = loss +  decay
     metrics['decay'] = decay
     metrics['raw_loss'] = loss
@@ -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)
+    loss, (state, metrics) = train_loss(apply, params, state, batch, index)  # pyrefly: ignore[bad-argument-type]
     for name, loss_fn in extra_losses.items():
       extra_loss, (unused_state,
-                   extra_metrics) = loss_fn(apply, params, state, batch, index)
+                   extra_metrics) = loss_fn(apply, params, state, batch, index)  # pyrefly: ignore[bad-argument-type]
       metrics[f'{name}:loss'] = extra_loss
       for key, value in extra_metrics.items():
         metrics[f'{name}:{key}'] = value
@@ -126,7 +126,7 @@ def combined_loss(enn: base.EpistemicNetwork[base.Input, base.Output],
       extra_loss, (unused_state,
                    extra_metrics) = loss_fn(enn, params, state, batch, key)
       metrics[f'{name}:loss'] = extra_loss
-      for key, value in extra_metrics.items():
+      for key, value in extra_metrics.items():  # pyrefly: ignore[bad-assignment]
         metrics[f'{name}:{key}'] = value
     return loss, (state, metrics)
 
@@ -171,7 +171,7 @@ def loss_fn(enn: base.EpistemicNetwork[base.Input, base.Output],
       for name, value in metrics.items():
         combined_metrics[f'{loss_config.name}:{name}'] = value
       combined_loss += loss_config.weight * loss
-    return combined_loss, (state, combined_metrics)
+    return combined_loss, (state, combined_metrics)  # pyrefly: ignore[bad-return]
 
   return loss_fn
 
@@ -196,11 +196,11 @@ def new_loss(
       batch: datasets.ArrayBatch,
       key: chex.PRNGKey
   ) -> base.LossOutput:
-    enn = networks.wrap_enn_as_enn_no_state(enn)
-    loss, metrics = loss_fn(enn, params, batch, key)
+    enn = networks.wrap_enn_as_enn_no_state(enn)  # pyrefly: ignore[bad-assignment]
+    loss, metrics = loss_fn(enn, params, batch, key)  # pyrefly: ignore[bad-argument-type]
     return loss, (constant_state, metrics)
 
-  return new_loss
+  return new_loss  # pyrefly: ignore[bad-return]
 
 
 def wrap_single_loss_no_state_as_single_loss(
@@ -223,10 +223,10 @@ def apply_no_state(params: hk.Params,
       output, unused_state = apply(params, constant_state, x, z)
       return output
 
-    loss, metrics = single_loss(apply_no_state, params, batch, index)
+    loss, metrics = single_loss(apply_no_state, params, batch, index)  # pyrefly: ignore[bad-argument-type]
     return loss, (constant_state, metrics)
 
-  return new_loss
+  return new_loss  # pyrefly: ignore[bad-return]
 
 
 def add_data_noise_no_state(
@@ -251,7 +251,7 @@ def add_l2_weight_decay_no_state(
 ) -> losses_base.LossFnNoState:
   """Adds scale * l2 weight decay to an existing loss function."""
   try:  # Scale is numeric.
-    scale = jnp.sqrt(scale)
+    scale = jnp.sqrt(scale)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     scale_fn = lambda ps: jax.tree_util.tree_map(lambda p: scale * p, ps)
   except TypeError:
     scale_fn = scale  # Assuming scale is a Callable.
@@ -261,7 +261,7 @@ def new_loss(enn: networks.EnnNoState,
                batch: datasets.ArrayBatch,
                key: chex.PRNGKey) -> losses_base.LossOutputNoState:
     loss, metrics = loss_fn(enn, params, batch, key)
-    decay = l2_weights_with_predicate(scale_fn(params), predicate)
+    decay = l2_weights_with_predicate(scale_fn(params), predicate)  # pyrefly: ignore[not-callable]
     total_loss = loss +  decay
     metrics['decay'] = decay
     metrics['raw_loss'] = loss
@@ -303,7 +303,7 @@ def combined_loss(enn: networks.EnnNoState,
     for name, loss_fn in extra_losses.items():
       extra_loss, extra_metrics = loss_fn(enn, params, batch, key)
       metrics[f'{name}:loss'] = extra_loss
-      for key, value in extra_metrics.items():
+      for key, value in extra_metrics.items():  # pyrefly: ignore[bad-assignment]
         metrics[f'{name}:{key}'] = value
     return loss, metrics
 
@@ -340,6 +340,6 @@ def loss_fn(enn: networks.EnnNoState,
       for name, value in metrics.items():
         combined_metrics[f'{loss_config.name}:{name}'] = value
       combined_loss += loss_config.weight * loss
-    return combined_loss, combined_metrics
+    return combined_loss, combined_metrics  # pyrefly: ignore[bad-return]
 
   return loss_fn
diff --git a/enn/losses/vi_losses.py b/enn/losses/vi_losses.py
index 6785a78..45d170d 100644
--- a/enn/losses/vi_losses.py
+++ b/enn/losses/vi_losses.py
@@ -45,7 +45,7 @@ def get_awgn_loglike_fn(
 
   def log_likelihood_fn(out: networks.Output, batch: datasets.ArrayBatch):
     chex.assert_shape(batch.y, (None, 1))
-    err_sq = jnp.mean(jnp.square(networks.parse_net_output(out) - batch.y))
+    err_sq = jnp.mean(jnp.square(networks.parse_net_output(out) - batch.y))  # pyrefly: ignore[unsupported-operation]
     return -0.5 * err_sq / sigma_w**2
 
   return log_likelihood_fn
@@ -68,7 +68,7 @@ def get_categorical_loglike_fn(
   def log_likelihood_fn(out: networks.Output, batch: datasets.ArrayBatch):
     chex.assert_shape(batch.y, (None, 1))
     logits = networks.parse_net_output(out)
-    labels = jax.nn.one_hot(batch.y[:, 0], num_classes)
+    labels = jax.nn.one_hot(batch.y[:, 0], num_classes)  # pyrefly: ignore[bad-index]
     return jnp.mean(
         jnp.sum(labels * jax.nn.log_softmax(logits), axis=1))
 
@@ -114,7 +114,7 @@ def sum_log_scale_mixture_normal(
 
 def normal_log_prob(latent: chex.Array, sigma: float = 1, mu: float = 0):
   """Compute un-normalized log probability of a normal RV."""
-  latent, _ = jax.tree.flatten(latent)
+  latent, _ = jax.tree.flatten(latent)  # pyrefly: ignore[bad-assignment]
   latent = jax.tree_util.tree_map(lambda x: x.flatten(), latent)
   latent = jnp.concatenate(latent)
   latent_dim = len(latent)
diff --git a/enn/metrics/calibration.py b/enn/metrics/calibration.py
index 0ed342c..9f411d2 100644
--- a/enn/metrics/calibration.py
+++ b/enn/metrics/calibration.py
@@ -77,7 +77,7 @@ def _get_init_stats(self,) -> metrics_base.MetricsState:
     return metrics_base.MetricsState(
         value=0,
         count=0,
-        extra=init_ece_stats,
+        extra=init_ece_stats,  # pyrefly: ignore[bad-argument-type]
     )
 
   def __call__(
diff --git a/enn/metrics/joint.py b/enn/metrics/joint.py
index 98d1acf..5a545bb 100644
--- a/enn/metrics/joint.py
+++ b/enn/metrics/joint.py
@@ -175,7 +175,7 @@ def reshape_to_smaller_batches(
 
   # 1.1. Discard extra data if needed.
   logits = logits[:, :num_data, :]
-  labels = labels[:num_data, :]
+  labels = labels[:num_data, :]  # pyrefly: ignore[bad-index]
   chex.assert_shape(logits, [num_enn_samples, num_data, num_classes])
   chex.assert_shape(labels, [num_data, 1])
 
diff --git a/enn/metrics/marginal.py b/enn/metrics/marginal.py
index ef5359e..b3fc4b2 100644
--- a/enn/metrics/marginal.py
+++ b/enn/metrics/marginal.py
@@ -24,7 +24,7 @@
 
 def make_nll_marginal_calculator() -> metrics_base.MetricCalculator:
   """Returns a MetricCalculator for marginal negative log likelihood (nll)."""
-  return lambda x, y: -1 * calculate_marginal_ll(x, y)
+  return lambda x, y: -1 * calculate_marginal_ll(x, y)  # pyrefly: ignore[bad-return]
 
 
 def make_accuracy_calculator() -> metrics_base.MetricCalculator:
@@ -62,6 +62,6 @@ def calculate_accuracy(logits: chex.Array, labels: chex.Array) -> float:
 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
-  assigned_probs = probs[jnp.arange(num_data), jnp.squeeze(labels)]
+  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
diff --git a/enn/metrics/metrics_test.py b/enn/metrics/metrics_test.py
index 31ee6b0..35de47c 100644
--- a/enn/metrics/metrics_test.py
+++ b/enn/metrics/metrics_test.py
@@ -137,7 +137,7 @@ def test_ece_calculator(
 
     # Check that ece results by our calculator and tfp calculator are the same.
     self.assertAlmostEqual(
-        our_ece_value, tfp_ece_value,
+        our_ece_value, tfp_ece_value,  # pyrefly: ignore[unbound-name]
         msg=f'our_ece_value={our_ece_value} not close enough to tfp_ece_value',
         delta=5e-2,
     )
diff --git a/enn/networks/bert/base.py b/enn/networks/bert/base.py
index 490f587..73dfd9b 100644
--- a/enn/networks/bert/base.py
+++ b/enn/networks/bert/base.py
@@ -127,5 +127,5 @@ def bert_large() -> BertConfig:
 
 class BertConfigs(enum.Enum):
   """Configs for BERT models."""
-  BERT_SMALL: BertConfig = bert_small()  # ~110M params
-  BERT_LARGE: BertConfig = bert_large()  # ~340M params
+  BERT_SMALL: BertConfig = bert_small()  # ~110M params  # pyrefly: ignore[invalid-annotation]
+  BERT_LARGE: BertConfig = bert_large()  # ~340M params  # pyrefly: ignore[invalid-annotation]
diff --git a/enn/networks/categorical_ensembles.py b/enn/networks/categorical_ensembles.py
index 353065b..609b2d8 100644
--- a/enn/networks/categorical_ensembles.py
+++ b/enn/networks/categorical_ensembles.py
@@ -46,7 +46,7 @@ def __init__(self, output_sizes: Sequence[int], atoms: chex.Array):
     super().__init__(name='categorical_regression_mlp')
     self.dim_out = output_sizes[-1]
     self.atoms = jnp.array(atoms)
-    self.output_sizes = list(output_sizes[:-1]) + [self.dim_out * len(atoms)]
+    self.output_sizes = list(output_sizes[:-1]) + [self.dim_out * len(atoms)]  # pyrefly: ignore[bad-argument-type]
 
   def __call__(self, inputs: chex.Array) -> chex.Array:
     """Apply MLP and wrap outputs appropriately."""
diff --git a/enn/networks/combiners.py b/enn/networks/combiners.py
index e13f404..73cf964 100644
--- a/enn/networks/combiners.py
+++ b/enn/networks/combiners.py
@@ -35,7 +35,7 @@ def combine_naive_enn(
     base_enn: enn_base.EpistemicNetwork[enn_base.Input, BaseOutput],
     parse_base_network: tp.Callable[
         [BaseOutput], HeadInput
-    ] = utils.parse_net_output,
+    ] = utils.parse_net_output,  # pyrefly: ignore[bad-function-definition]
 ) -> enn_base.EpistemicNetwork[enn_base.Input, networks_base.Output]:
   """Combines a base enn and a head enn naively without optimization.
 
@@ -101,7 +101,7 @@ def init(key: chex.PRNGKey,
     return (params, state)
 
   return enn_base.EpistemicNetwork[enn_base.Input, networks_base.Output](
-      apply, init, base_enn.indexer
+      apply, init, base_enn.indexer  # pyrefly: ignore[bad-argument-type]
   )
 
 
@@ -112,7 +112,7 @@ def make_optimized_forward(
     key: chex.PRNGKey,
     parse_base_network: tp.Callable[
         [BaseOutput], HeadInput
-    ] = utils.parse_net_output,
+    ] = utils.parse_net_output,  # pyrefly: ignore[bad-function-definition]
 ) -> forwarders.EnnBatchFwd[enn_base.Input]:
   """Combines base enn and head enn for multiple ENN samples.
 
diff --git a/enn/networks/dropout.py b/enn/networks/dropout.py
index 17079bd..49738b9 100644
--- a/enn/networks/dropout.py
+++ b/enn/networks/dropout.py
@@ -124,7 +124,7 @@ def apply(params: hk.Params, x: chex.Array,
       net_out = transformed.apply(params, x, z)
       return net_out
 
-    apply = network_utils.wrap_apply_no_state_as_apply(apply)
+    apply = network_utils.wrap_apply_no_state_as_apply(apply)  # pyrefly: ignore[bad-argument-type]
     init = network_utils.wrap_init_no_state_as_init(transformed.init)
 
     super().__init__(apply, init, indexer)
diff --git a/enn/networks/einsum_mlp.py b/enn/networks/einsum_mlp.py
index 0b6a7a2..2fc091d 100644
--- a/enn/networks/einsum_mlp.py
+++ b/enn/networks/einsum_mlp.py
@@ -70,8 +70,8 @@ def init(key: chex.PRNGKey, x: chex.Array,
   indexer = indexers.EnsembleIndexer(num_ensemble)
 
   # TODO(author3): Change apply and init fns above to work with state.
-  apply = network_utils.wrap_apply_no_state_as_apply(apply)
-  init = network_utils.wrap_init_no_state_as_init(init)
+  apply = network_utils.wrap_apply_no_state_as_apply(apply)  # pyrefly: ignore[bad-argument-type]
+  init = network_utils.wrap_init_no_state_as_init(init)  # pyrefly: ignore[bad-argument-type]
   return networks_base.EnnArray(apply, init, indexer)
 
 
@@ -111,10 +111,10 @@ def apply_with_prior(
     ensemble_train, state = enn.apply(params, state, x, z)
     ensemble_prior, _ = enn.apply(prior_params, prior_state, x, z)
     output = networks_base.OutputWithPrior(
-        train=ensemble_train, prior=ensemble_prior * prior_scale)
+        train=ensemble_train, prior=ensemble_prior * prior_scale)  # pyrefly: ignore[bad-argument-type, unsupported-operation]
     return output, state
 
-  return networks_base.EnnArray(apply_with_prior, enn.init, enn.indexer)
+  return networks_base.EnnArray(apply_with_prior, enn.init, enn.indexer)  # pyrefly: ignore[bad-argument-type]
 
 
 # TODO(author3): Come up with a better name and use ensembles.py instead.
diff --git a/enn/networks/ensembles.py b/enn/networks/ensembles.py
index fc2aa2e..6443e27 100644
--- a/enn/networks/ensembles.py
+++ b/enn/networks/ensembles.py
@@ -52,7 +52,7 @@ def apply(params: hk.Params, inputs: chex.Array,
       return model.apply(sub_params, inputs)
 
     indexer = indexers.EnsembleIndexer(num_ensemble)
-    super().__init__(apply, init, indexer)
+    super().__init__(apply, init, indexer)  # pyrefly: ignore[bad-argument-type]
 
 
 class EnsembleWithState(networks_base.EnnArray):
@@ -86,7 +86,7 @@ def apply(params: hk.Params, states: hk.State, inputs: chex.Array,
       return out, new_states
 
     indexer = indexers.EnsembleIndexer(num_ensemble)
-    super().__init__(apply, init, indexer)
+    super().__init__(apply, init, indexer)  # pyrefly: ignore[bad-argument-type]
 
 
 def make_mlp_ensemble_prior_fns(
diff --git a/enn/networks/epinet/base.py b/enn/networks/epinet/base.py
index 43504af..de178ea 100644
--- a/enn/networks/epinet/base.py
+++ b/enn/networks/epinet/base.py
@@ -135,4 +135,4 @@ def init(key: chex.PRNGKey,
     else:
       return {**params, **base_params_init}, {**state, **base_state_init}
 
-  return networks_base.EnnArray(apply, init, epinet.indexer)
+  return networks_base.EnnArray(apply, init, epinet.indexer)  # pyrefly: ignore[bad-argument-type]
diff --git a/enn/networks/epinet/priors.py b/enn/networks/epinet/priors.py
index f2fbe3e..ccd5b63 100644
--- a/enn/networks/epinet/priors.py
+++ b/enn/networks/epinet/priors.py
@@ -124,7 +124,7 @@ def conv_net(x):
 
   def prior_fn(x: chex.Array, z: chex.Array) -> chex.Array:
     out = [ensemble.apply(params, x, index) for index in range(num_ensemble)]
-    out = jnp.stack(out, axis=-1)
+    out = jnp.stack(out, axis=-1)  # pyrefly: ignore[bad-argument-type]
     return jnp.dot(out, z)
 
   return jax.jit(prior_fn)
@@ -141,14 +141,14 @@ def make_cifar_conv_prior(
   rng = hk.PRNGSequence(seed)
 
   if kernel_shapes is None:
-    kernel_shapes = [[5, 5]] * len(output_channels)
-  assert len(output_channels) == len(kernel_shapes)
+    kernel_shapes = [[5, 5]] * len(output_channels)  # pyrefly: ignore[bad-assignment]
+  assert len(output_channels) == len(kernel_shapes)  # pyrefly: ignore[bad-argument-type]
 
   def conv_net(x):
     x = jax.image.resize(x, [x.shape[0], 32, 32, 3], method='bilinear')
     for i, channels in enumerate(output_channels):
       x = hk.Conv2D(output_channels=channels,
-                    kernel_shape=kernel_shapes[i],
+                    kernel_shape=kernel_shapes[i],  # pyrefly: ignore[unsupported-operation]
                     stride=2, name='prior_conv')(x)
       x = jax.nn.relu(x)
     x = hk.Flatten()(x)
@@ -163,7 +163,7 @@ def conv_net(x):
 
   def prior_fn(x: chex.Array, z: chex.Array) -> chex.Array:
     out = [ensemble.apply(params, x, index) for index in range(num_ensemble)]
-    out = jnp.stack(out, axis=-1)
+    out = jnp.stack(out, axis=-1)  # pyrefly: ignore[bad-argument-type]
     return jnp.dot(out, z)
 
   return jax.jit(prior_fn)
diff --git a/enn/networks/gaussian_enn.py b/enn/networks/gaussian_enn.py
index 09a9370..fbe9fbd 100644
--- a/enn/networks/gaussian_enn.py
+++ b/enn/networks/gaussian_enn.py
@@ -89,8 +89,8 @@ def net_fn(inputs: chex.Array) -> chex.Array:
     init = lambda rng, x, z: transformed.init(rng, x)
 
     # TODO(author3): Change apply and init fns above to work with state.
-    apply = network_utils.wrap_apply_no_state_as_apply(apply)
-    init = network_utils.wrap_init_no_state_as_init(init)
+    apply = network_utils.wrap_apply_no_state_as_apply(apply)  # pyrefly: ignore[bad-argument-type]
+    init = network_utils.wrap_init_no_state_as_init(init)  # pyrefly: ignore[bad-argument-type]
 
     super().__init__(apply, init, indexer=indexers.PrngIndexer(),)
 
diff --git a/enn/networks/mlp.py b/enn/networks/mlp.py
index db1ef10..0611f69 100644
--- a/enn/networks/mlp.py
+++ b/enn/networks/mlp.py
@@ -63,7 +63,7 @@ def __call__(self, inputs: chex.Array) -> networks_base.OutputWithPrior:
     exposed_features = [inputs]
     for i, layer_feature in enumerate(layers_features):
       # Add this layer feature if the expose flag for this layer is True
-      if self.expose_layers[i]:
+      if self.expose_layers[i]:  # pyrefly: ignore[unsupported-operation]
         exposed_features.append(layer_feature)
 
     exposed_features = jnp.concatenate(exposed_features, axis=1)
@@ -73,7 +73,7 @@ def __call__(self, inputs: chex.Array) -> networks_base.OutputWithPrior:
     return networks_base.OutputWithPrior(
         train=out,
         prior=jnp.zeros_like(out),
-        extra=extra
+        extra=extra  # pyrefly: ignore[bad-argument-type]
     )
 
 
diff --git a/enn/networks/priors.py b/enn/networks/priors.py
index a8ac130..ec75958 100644
--- a/enn/networks/priors.py
+++ b/enn/networks/priors.py
@@ -136,13 +136,13 @@ def make_random_feat_gp(
   bias = 2 * jnp.pi * jax.random.uniform(
       bias_key, shape=[1, num_feat, output_dim])
   alpha = jax.random.normal(alpha_key, shape=[num_feat]) / jnp.sqrt(num_feat)
-  gamma = _parse_gamma(gamma, num_feat, gamma_key)
+  gamma = _parse_gamma(gamma, num_feat, gamma_key)  # pyrefly: ignore[bad-assignment]
 
   def gp_instance(inputs: chex.Array) -> chex.Array:
     """Assumes one batch dimension and flattens input to match that."""
     flat_inputs = jax.vmap(jnp.ravel)(inputs)
     input_embedding = jnp.einsum('bi,kio->bko', flat_inputs, weights)
-    random_feats = jnp.cos(gamma * input_embedding + bias)
+    random_feats = jnp.cos(gamma * input_embedding + bias)  # pyrefly: ignore[unsupported-operation]
     return scale * jnp.einsum('bko,k->bo', random_feats, alpha)
 
   return gp_instance
diff --git a/enn/networks/resnet/lib.py b/enn/networks/resnet/lib.py
index 6a58715..f693a79 100644
--- a/enn/networks/resnet/lib.py
+++ b/enn/networks/resnet/lib.py
@@ -265,56 +265,56 @@ def _make_resnet_blocks(config: ResNetConfig) -> Sequence[ResBlock]:
 
 class CanonicalResNets(enum.Enum):
   """Canonical ResNet configs."""
-  RESNET_18: ResNetConfig = ResNetConfig(
+  RESNET_18: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(16, 32, 64),
       blocks_per_group=(2, 2, 2),
       strides_per_group=(1, 2, 2),
       resnet_block_version='V1',
   )
-  RESNET_32: ResNetConfig = ResNetConfig(
+  RESNET_32: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(16, 32, 64),
       blocks_per_group=(5, 5, 5),
       strides_per_group=(1, 2, 2),
       resnet_block_version='V1',
   )
-  RESNET_44: ResNetConfig = ResNetConfig(
+  RESNET_44: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(16, 32, 64),
       blocks_per_group=(7, 7, 7),
       strides_per_group=(1, 2, 2),
       resnet_block_version='V1',
   )
-  RESNET_56: ResNetConfig = ResNetConfig(
+  RESNET_56: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(16, 32, 64),
       blocks_per_group=(9, 9, 9),
       strides_per_group=(1, 2, 2),
       resnet_block_version='V1',
   )
-  RESNET_110: ResNetConfig = ResNetConfig(
+  RESNET_110: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(16, 32, 64),
       blocks_per_group=(18, 18, 18),
       strides_per_group=(1, 2, 2),
       resnet_block_version='V1',
   )
 
-  RESNET_50: ResNetConfig = ResNetConfig(
+  RESNET_50: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(256, 512, 1024, 2048),
       blocks_per_group=(3, 4, 6, 3),
       strides_per_group=(1, 2, 2, 2),
       resnet_block_version='V2',
   )
-  RESNET_101: ResNetConfig = ResNetConfig(
+  RESNET_101: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(256, 512, 1024, 2048),
       blocks_per_group=(3, 4, 23, 3),
       strides_per_group=(1, 2, 2, 2),
       resnet_block_version='V2',
   )
-  RESNET_152: ResNetConfig = ResNetConfig(
+  RESNET_152: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(256, 512, 1024, 2048),
       blocks_per_group=(3, 8, 36, 3),
       strides_per_group=(1, 2, 2, 2),
       resnet_block_version='V2',
   )
-  RESNET_200: ResNetConfig = ResNetConfig(
+  RESNET_200: ResNetConfig = ResNetConfig(  # pyrefly: ignore[invalid-annotation]
       channels_per_group=(256, 512, 1024, 2048),
       blocks_per_group=(3, 24, 36, 3),
       strides_per_group=(1, 2, 2, 2),
diff --git a/enn/networks/utils.py b/enn/networks/utils.py
index edf8172..38f3f09 100644
--- a/enn/networks/utils.py
+++ b/enn/networks/utils.py
@@ -72,16 +72,16 @@ def wrap_net_fn_as_enn(
   def apply(
       params: hk.Params,
       state: hk.State,
-      inputs: base.Input,
+      inputs: base.Input,  # pyrefly: ignore[invalid-type-var]
       index: base.Index,
-  ) -> Tuple[base.Output, hk.State]:
+  ) -> Tuple[base.Output, hk.State]:  # pyrefly: ignore[invalid-type-var]
     del index
     return transformed.apply(params, state, inputs)
 
-  return base.EpistemicNetwork[base.Input, base.Output](
+  return base.EpistemicNetwork[base.Input, base.Output](  # pyrefly: ignore[bad-return]
       apply=apply,
-      init=lambda k, x, z: transformed.init(k, x),
-      indexer=lambda k: k,
+      init=lambda k, x, z: transformed.init(k, x),  # pyrefly: ignore[bad-argument-type]
+      indexer=lambda k: k,  # pyrefly: ignore[bad-argument-type]
   )
 
 
@@ -89,8 +89,8 @@ def wrap_transformed_as_enn_no_state(
     transformed: hk.Transformed) -> networks_base.EnnNoState:
   """Wraps a simple transformed function y = f(x) as an ENN."""
   return networks_base.EnnNoState(
-      apply=lambda params, x, z: transformed.apply(params, x),
-      init=lambda key, x, z: transformed.init(key, x),
+      apply=lambda params, x, z: transformed.apply(params, x),  # pyrefly: ignore[bad-argument-type]
+      init=lambda key, x, z: transformed.init(key, x),  # pyrefly: ignore[bad-argument-type]
       indexer=lambda key: key,
   )
 
@@ -100,9 +100,9 @@ def wrap_transformed_as_enn(
 ) -> networks_base.EnnArray:
   """Wraps a simple transformed function y = f(x) as an ENN."""
   apply = lambda params, x, z: transformed.apply(params, x)
-  apply = wrap_apply_no_state_as_apply(apply)
+  apply = wrap_apply_no_state_as_apply(apply)  # pyrefly: ignore[bad-argument-type]
   init = lambda key, x, z: transformed.init(key, x)
-  init = wrap_init_no_state_as_init(init)
+  init = wrap_init_no_state_as_init(init)  # pyrefly: ignore[bad-argument-type]
   return networks_base.EnnArray(
       apply=apply,
       init=init,
@@ -140,8 +140,8 @@ def apply(
     return output
 
   return networks_base.EnnNoState(
-      apply=apply,
-      init=init,
+      apply=apply,  # pyrefly: ignore[bad-argument-type]
+      init=init,  # pyrefly: ignore[bad-argument-type]
       indexer=enn.indexer,
   )
 
@@ -156,7 +156,7 @@ def new_apply(
       index: base.Index,
   ) -> Tuple[networks_base.Output, hk.State]:
     return (apply(params, inputs, index), {})
-  return new_apply
+  return new_apply  # pyrefly: ignore[bad-return]
 
 
 def wrap_init_no_state_as_init(
@@ -169,7 +169,7 @@ def new_init(
       index: base.Index,
   ) -> Tuple[hk.Params, hk.State]:
     return (init(key, inputs, index), {})
-  return new_init
+  return new_init  # pyrefly: ignore[bad-return]
 
 
 def scale_enn_output(
@@ -210,7 +210,7 @@ def centered_apply(params: hk.Params, x: chex.Array,
     normalized_x = (x - x_mean) / (x_std + 1e-9)
     return enn.apply(params, normalized_x, z)
 
-  return networks_base.EnnNoState(centered_apply, enn.init, enn.indexer)
+  return networks_base.EnnNoState(centered_apply, enn.init, enn.indexer)  # pyrefly: ignore[bad-argument-type]
 
 
 def make_centered_enn(
diff --git a/enn/supervised/multiloss_experiment.py b/enn/supervised/multiloss_experiment.py
index e48e9e6..86bfa65 100644
--- a/enn/supervised/multiloss_experiment.py
+++ b/enn/supervised/multiloss_experiment.py
@@ -96,7 +96,7 @@ def forward(params: hk.Params, state: hk.State, inputs: chex.Array,
                 key: chex.PRNGKey) -> chex.Array:
       index = self.enn.indexer(key)
       out, _ = self.enn.apply(params, state, inputs, index)
-      return out
+      return out  # pyrefly: ignore[bad-return]
     self._forward = jax.jit(forward)
 
     # Define the SGD step on the loss
@@ -116,7 +116,7 @@ def sgd_step(
       )
       new_params = optax.apply_updates(state.params, updates)
       new_state = TrainingState(
-          params=new_params,
+          params=new_params,  # pyrefly: ignore[bad-argument-type]
           network_state=network_state,
           opt_state=new_opt_state,
       )
@@ -164,7 +164,7 @@ def train(self, num_batches: int):
                 next(dataset),
                 next(self.rng),
             )
-            metrics.update({
+            metrics.update({  # pyrefly: ignore[no-matching-overload]
                 'dataset': name,
                 'step': self.step,
                 'sgd': False,
diff --git a/enn/supervised/sgd_experiment.py b/enn/supervised/sgd_experiment.py
index 4ecd3a1..51fa339 100644
--- a/enn/supervised/sgd_experiment.py
+++ b/enn/supervised/sgd_experiment.py
@@ -103,7 +103,7 @@ def forward(params: hk.Params,
                 key: chex.PRNGKey) -> chex.Array:
       index = self.enn.indexer(key)
       out, unused_state = self.enn.apply(params, state, inputs, index)
-      return out
+      return out  # pyrefly: ignore[bad-return]
     self._forward = jax.jit(forward)
 
     # Batched forward at multiple random indices
@@ -123,7 +123,7 @@ def sgd_step(
       updates, new_opt_state = optimizer.update(grads, training_state.opt_state)
       new_params = optax.apply_updates(training_state.params, updates)
       new_state = TrainingState(
-          params=new_params,
+          params=new_params,  # pyrefly: ignore[bad-argument-type]
           network_state=network_state,
           opt_state=new_opt_state,
       )
@@ -166,7 +166,7 @@ def train(self, num_batches: int):
 
       # Periodically evaluate the other datasets.
       if self._should_eval and self.step % self._eval_log_freq == 0:
-        for name, dataset in self._eval_datasets.items():
+        for name, dataset in self._eval_datasets.items():  # pyrefly: ignore[missing-attribute]
           # Evaluation happens on a single batch
           eval_batch = next(dataset)
           eval_metrics = {'dataset': name, 'step': self.step, 'sgd': False}
@@ -178,7 +178,7 @@ def train(self, num_batches: int):
               jax.random.split(next(self.rng), self._eval_enn_samples),
           )
           logits = networks.parse_net_output(net_out)
-          for metric_name, metric_calc in self._eval_metrics.items():
+          for metric_name, metric_calc in self._eval_metrics.items():  # pyrefly: ignore[missing-attribute]
             eval_metrics.update({
                 metric_name: metric_calc(logits, eval_batch.y),
             })
diff --git a/enn/utils.py b/enn/utils.py
index e9545f2..4758282 100644
--- a/enn/utils.py
+++ b/enn/utils.py
@@ -53,11 +53,11 @@ def _clean_batch_data(data: ds_base.ArrayBatch) -> ds_base.ArrayBatch:
 
   # Data index to identify each instance
   if data.data_index is None:
-    data = dataclasses.replace(data, data_index=np.arange(len(data.y))[:, None])
+    data = dataclasses.replace(data, data_index=np.arange(len(data.y))[:, None])  # pyrefly: ignore[bad-argument-type]
 
   # Weights to say how much each data.point is work
   if data.weights is None:
-    data = dataclasses.replace(data, weights=np.ones(len(data.y))[:, None])
+    data = dataclasses.replace(data, weights=np.ones(len(data.y))[:, None])  # pyrefly: ignore[bad-argument-type]
   return data
 
 
@@ -66,11 +66,11 @@ def make_batch_iterator(data: ds_base.ArrayBatch,
                         seed: int = 0) -> ds_base.ArrayBatchIterator:
   """Converts toy-like training data to batch_iterator for sgd training."""
   data = _clean_batch_data(data)
-  n_data = len(data.y)
+  n_data = len(data.y)  # pyrefly: ignore[bad-argument-type]
   if not batch_size:
     batch_size = n_data
 
-  ds = tf.data.Dataset.from_tensor_slices(data).cache()
+  ds = tf.data.Dataset.from_tensor_slices(data).cache()  # pyrefly: ignore[bad-argument-type]
   ds = ds.shuffle(min(n_data, 50 * batch_size), seed=seed)
   ds = ds.repeat().batch(batch_size)