From d79d6a54d0cd0be4bffb9b3d5101e409fcf3b253 Mon Sep 17 00:00:00 2001
From: Oleh Prypin <oprypin@google.com>
Date: Mon, 5 Oct 2026 08:21:15 -0700
Subject: [PATCH] Replace `# pytype: disable` suppressions with `# pyrefly:
 ignore`

PiperOrigin-RevId: 993699001
---
 edward2/jax/nn/attention.py                   |  2 +-
 edward2/jax/nn/random_feature.py              |  8 +--
 edward2/tensorflow/initializers.py            |  2 +-
 .../layers/bayesian_linear_model.py           |  4 +-
 edward2/tensorflow/layers/convolutional.py    | 28 ++++-----
 edward2/tensorflow/layers/dense.py            | 20 +++----
 edward2/tensorflow/layers/discrete_flows.py   |  6 +-
 edward2/tensorflow/layers/embeddings.py       |  2 +-
 edward2/tensorflow/layers/gaussian_process.py |  4 +-
 edward2/tensorflow/layers/heteroscedastic.py  | 58 +++++++++----------
 edward2/tensorflow/layers/hetsngp.py          |  6 +-
 edward2/tensorflow/layers/made.py             |  4 +-
 edward2/tensorflow/layers/neural_process.py   | 10 ++--
 edward2/tensorflow/layers/normalization.py    |  6 +-
 .../tensorflow/layers/posterior_network.py    |  4 +-
 edward2/tensorflow/layers/recurrent.py        | 12 ++--
 .../tensorflow/layers/stochastic_output.py    |  2 +-
 edward2/tensorflow/layers/utils.py            |  8 +--
 edward2/tensorflow/random_variable.py         |  2 +-
 .../no_u_turn_sampler/logistic_regression.py  |  4 +-
 examples/no_u_turn_sampler/nuts.py            |  6 +-
 21 files changed, 99 insertions(+), 99 deletions(-)

diff --git a/edward2/jax/nn/attention.py b/edward2/jax/nn/attention.py
index b1e6ccb9..7990c261 100644
--- a/edward2/jax/nn/attention.py
+++ b/edward2/jax/nn/attention.py
@@ -188,7 +188,7 @@ def f(x, *args, **kwargs):
         broadcast_dropout=self.broadcast_dropout,  # pyrefly: ignore[unexpected-keyword]
         deterministic=deterministic,  # pyrefly: ignore[unexpected-keyword]
         dtype=self.dtype,  # pyrefly: ignore[unexpected-keyword]
-        precision=self.precision)  # pytype: disable=wrong-keyword-args
+        precision=self.precision)  # pyrefly: ignore[unexpected-keyword]
     # back to the original inputs dimensions
     def dense_fn2(name):
       dense_lyr = dense.DenseBatchEnsemble(
diff --git a/edward2/jax/nn/random_feature.py b/edward2/jax/nn/random_feature.py
index 64670762..e125a5cb 100644
--- a/edward2/jax/nn/random_feature.py
+++ b/edward2/jax/nn/random_feature.py
@@ -136,8 +136,8 @@ def __call__(self,
 
     # Returns predictive logits, covmat and (optionally) random features.
     if return_random_features:
-      return gp_logits, gp_covmat, gp_features  # pytype: disable=bad-return-type  # jax-ndarray
-    return gp_logits, gp_covmat  # pytype: disable=bad-return-type  # jax-ndarray
+      return gp_logits, gp_covmat, gp_features  # pyrefly: ignore[bad-return]
+    return gp_logits, gp_covmat  # pyrefly: ignore[bad-return]
 
 
 class RandomFourierFeatures(nn.Module):
@@ -641,7 +641,7 @@ def _compute_mc_samples(self, locs: Array, scale: Array,
     """
     locs = jnp.expand_dims(locs, axis=1)
 
-    noise_samples = self._compute_noise_samples(scale, num_samples)  # pytype: disable=wrong-arg-types  # jax-ndarray
+    noise_samples = self._compute_noise_samples(scale, num_samples)  # pyrefly: ignore[bad-argument-type]
 
     latents = locs + noise_samples
     samples = jax.nn.sigmoid(latents / self.temperature)
@@ -677,7 +677,7 @@ def __call__(self, inputs: Array, training: int = True) -> Union[
     else:
       total_mc_samples = self.test_mc_samples
 
-    probs_mean = self._compute_mc_samples(locs, scale, total_mc_samples)  # pytype: disable=wrong-arg-types  # jax-ndarray
+    probs_mean = self._compute_mc_samples(locs, scale, total_mc_samples)  # pyrefly: ignore[bad-argument-type]
 
     probs_mean = jnp.clip(probs_mean, min=self.eps)
     log_probs = jnp.log(probs_mean)
diff --git a/edward2/tensorflow/initializers.py b/edward2/tensorflow/initializers.py
index 8ccdfd20..1f7f02c0 100644
--- a/edward2/tensorflow/initializers.py
+++ b/edward2/tensorflow/initializers.py
@@ -781,7 +781,7 @@ def __init__(self, stddev=1.0, random_norm=True, seed=None):
     self.random_norm = random_norm
 
   def _sample_orthogonal_matrix(self, shape, dtype):
-    return super(OrthogonalRandomFeatures, self).__call__(shape, dtype=dtype)  # pytype: disable=attribute-error  # typed-keras
+    return super(OrthogonalRandomFeatures, self).__call__(shape, dtype=dtype)
 
   def __call__(self, shape, dtype=tf.float32, **kwargs):
     # Sample orthogonal matrices.
diff --git a/edward2/tensorflow/layers/bayesian_linear_model.py b/edward2/tensorflow/layers/bayesian_linear_model.py
index d6e36222..02e3c64c 100644
--- a/edward2/tensorflow/layers/bayesian_linear_model.py
+++ b/edward2/tensorflow/layers/bayesian_linear_model.py
@@ -42,7 +42,7 @@ def __init__(self, noise_variance, **kwargs):
     self.coeffs_precision_tril_op = None
     self.coeffs_mean = None
 
-  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def call(self, inputs):  # pyrefly: ignore[bad-override]
     if self.coeffs_mean is None and self.coeffs_precision_tril_op is None:
       # p(mean(ynew) | xnew) = Normal(ynew | mean = 0, variance = xnew xnew^T)
       predictive_mean = 0.
@@ -61,7 +61,7 @@ def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-para
     return generated_random_variables.Normal(loc=predictive_mean,  # pyrefly: ignore[missing-attribute]
                                              scale=tf.sqrt(predictive_variance))
 
-  def fit(self, x=None, y=None):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def fit(self, x=None, y=None):  # pyrefly: ignore[bad-override]
     # p(coeffs | x, y) = Normal(coeffs |
     #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
     #   covariance = (1/noise_variance x^T x + I)^{-1})
diff --git a/edward2/tensorflow/layers/convolutional.py b/edward2/tensorflow/layers/convolutional.py
index 4558583b..c0523d3a 100644
--- a/edward2/tensorflow/layers/convolutional.py
+++ b/edward2/tensorflow/layers/convolutional.py
@@ -86,9 +86,9 @@ def __init__(self,
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)  # pyrefly: ignore[not-callable]
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)  # pyrefly: ignore[missing-attribute, not-callable]
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)  # pyrefly: ignore[missing-attribute]
 
   def call(self, *args, **kwargs):
     self.call_weights()
@@ -154,9 +154,9 @@ def __init__(self,
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)  # pyrefly: ignore[not-callable]
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)  # pyrefly: ignore[missing-attribute, not-callable]
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)  # pyrefly: ignore[missing-attribute]
 
   def call(self, *args, **kwargs):
     self.call_weights()
@@ -434,10 +434,10 @@ def build(self, input_shape):
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.local_scale_initializer, tf.keras.layers.Layer):
-      self.local_scale = self.local_scale_initializer(self.local_scale.shape,  # pyrefly: ignore[not-callable]
+      self.local_scale = self.local_scale_initializer(self.local_scale.shape,
                                                       self.dtype)
     if isinstance(self.global_scale_initializer, tf.keras.layers.Layer):
-      self.global_scale = self.global_scale_initializer(self.global_scale.shape,  # pyrefly: ignore[not-callable]
+      self.global_scale = self.global_scale_initializer(self.global_scale.shape,
                                                         self.dtype)
     super().call_weights()
 
@@ -1448,7 +1448,7 @@ def build(self, input_shape):
 
     self.built = True
 
-  def call(self, inputs, routing_weights):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def call(self, inputs, routing_weights):  # pyrefly: ignore[bad-override]
     # Compute example dependent kernels
     kernels = tf.matmul(routing_weights, self.condconv_kernel)
     batch_size = self.batch_size
@@ -1665,7 +1665,7 @@ def build(self, input_shape):
         ndim=4, axes={channel_axis: input_dim})
     self.built = True
 
-  def call(self, inputs, routing_weights):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def call(self, inputs, routing_weights):  # pyrefly: ignore[bad-override]
     # Compute example dependent depthwise kernels
     depthwise_kernels = tf.matmul(routing_weights,
                                   self.depthwise_condconv_kernel)
@@ -1974,7 +1974,7 @@ def call(self, inputs):
     # Sample parameters for each example.
     if isinstance(self.alpha_initializer, tf.keras.layers.Layer):
       alpha = tf.clip_by_value(
-          self.alpha_initializer(  # pyrefly: ignore[not-callable]
+          self.alpha_initializer(
               self.alpha_shape,
               self.dtype).distribution.sample(examples_per_model),
           self.min_perturbation_value,
@@ -1984,7 +1984,7 @@ def call(self, inputs):
       alpha = tf.tile(self.alpha, [1, examples_per_model])
     if isinstance(self.gamma_initializer, tf.keras.layers.Layer):
       gamma = tf.clip_by_value(
-          self.gamma_initializer(  # pyrefly: ignore[not-callable]
+          self.gamma_initializer(
               self.gamma_shape,
               self.dtype).distribution.sample(examples_per_model),
           self.min_perturbation_value,
@@ -2005,7 +2005,7 @@ def call(self, inputs):
 
     if self.use_ensemble_bias:
       if isinstance(self.ensemble_bias_initializer, tf.keras.layers.Layer):
-        bias = self.ensemble_bias_initializer(  # pyrefly: ignore[not-callable]
+        bias = self.ensemble_bias_initializer(
             self.ensemble_bias_shape,
             self.dtype).distribution.sample(examples_per_model)
         bias = tf.transpose(bias, [1, 0, 2])
@@ -2179,7 +2179,7 @@ def call(self, inputs):
     # Sample parameters for each example.
     if isinstance(self.alpha_initializer, tf.keras.layers.Layer):
       alpha = tf.clip_by_value(
-          self.alpha_initializer(  # pyrefly: ignore[not-callable]
+          self.alpha_initializer(
               self.alpha_shape,
               self.dtype).distribution.sample(examples_per_model),
           self.min_perturbation_value,
@@ -2189,7 +2189,7 @@ def call(self, inputs):
       alpha = tf.tile(self.alpha, [1, examples_per_model])
     if isinstance(self.gamma_initializer, tf.keras.layers.Layer):
       gamma = tf.clip_by_value(
-          self.gamma_initializer(  # pyrefly: ignore[not-callable]
+          self.gamma_initializer(
               self.gamma_shape,
               self.dtype).distribution.sample(examples_per_model),
           self.min_perturbation_value,
@@ -2212,7 +2212,7 @@ def call(self, inputs):
 
     if self.use_ensemble_bias:
       if isinstance(self.ensemble_bias_initializer, tf.keras.layers.Layer):
-        bias = self.ensemble_bias_initializer(  # pyrefly: ignore[not-callable]
+        bias = self.ensemble_bias_initializer(
             self.ensemble_bias_shape,
             self.dtype).distribution.sample(examples_per_model)
         bias = tf.transpose(bias, [1, 0, 2])
diff --git a/edward2/tensorflow/layers/dense.py b/edward2/tensorflow/layers/dense.py
index 12472721..6b503279 100644
--- a/edward2/tensorflow/layers/dense.py
+++ b/edward2/tensorflow/layers/dense.py
@@ -73,9 +73,9 @@ def __init__(self,
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)  # pyrefly: ignore[not-callable]
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)  # pyrefly: ignore[not-callable]
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
 
   def call(self, *args, **kwargs):
     self.call_weights()
@@ -452,10 +452,10 @@ def build(self, input_shape):
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.local_scale_initializer, tf.keras.layers.Layer):
-      self.local_scale = self.local_scale_initializer(self.local_scale.shape,  # pyrefly: ignore[not-callable]
+      self.local_scale = self.local_scale_initializer(self.local_scale.shape,
                                                       self.dtype)
     if isinstance(self.global_scale_initializer, tf.keras.layers.Layer):
-      self.global_scale = self.global_scale_initializer(self.global_scale.shape,  # pyrefly: ignore[not-callable]
+      self.global_scale = self.global_scale_initializer(self.global_scale.shape,
                                                         self.dtype)
     super().call_weights()
 
@@ -808,8 +808,8 @@ def call(self, inputs):
     data, lambdas, e = inputs
     e1, e2 = e[:, :self.units], e[:, self.units:]
 
-    output = self.dense(data)  # pyrefly: ignore[not-callable]
-    delta_kernel = self.delta_dense(data) * e1  # pyrefly: ignore[not-callable]
+    output = self.dense(data)
+    delta_kernel = self.delta_dense(data) * e1
     output += delta_kernel
 
     batch_size = tf.shape(data)[0]
@@ -1104,7 +1104,7 @@ def call(self, inputs):
     # Sample parameters for each example.
     if isinstance(self.alpha_initializer, tf.keras.layers.Layer):
       alpha = tf.clip_by_value(
-          self.alpha_initializer(  # pyrefly: ignore[not-callable]
+          self.alpha_initializer(
               self.alpha_shape,
               self.dtype).distribution.sample(examples_per_model),
           self.min_perturbation_value,
@@ -1114,7 +1114,7 @@ def call(self, inputs):
       alpha = tf.expand_dims(self.alpha, 1)
     if isinstance(self.gamma_initializer, tf.keras.layers.Layer):
       gamma = tf.clip_by_value(
-          self.gamma_initializer(  # pyrefly: ignore[not-callable]
+          self.gamma_initializer(
               self.gamma_shape,
               self.dtype).distribution.sample(examples_per_model),
           self.min_perturbation_value,
@@ -1130,7 +1130,7 @@ def call(self, inputs):
 
     if self.use_ensemble_bias:
       if isinstance(self.ensemble_bias_initializer, tf.keras.layers.Layer):
-        bias = self.ensemble_bias_initializer(  # pyrefly: ignore[not-callable]
+        bias = self.ensemble_bias_initializer(
             self.ensemble_bias_shape,
             self.dtype).distribution.sample(examples_per_model)
         bias = tf.transpose(bias, [1, 0, 2])
@@ -1262,7 +1262,7 @@ def build(self, input_shape):
 
     self.built = True
 
-  def call(self, inputs, routing_weights):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def call(self, inputs, routing_weights):  # pyrefly: ignore[bad-override]
     # Compute example dependent kernels
     inputs = tf.expand_dims(inputs, 1)  # shape = [batch_size, 1, input_dim]
     # routing_weights is of shape [batch_size, num_experts]
diff --git a/edward2/tensorflow/layers/discrete_flows.py b/edward2/tensorflow/layers/discrete_flows.py
index 4d277a36..a8fcbe04 100644
--- a/edward2/tensorflow/layers/discrete_flows.py
+++ b/edward2/tensorflow/layers/discrete_flows.py
@@ -108,7 +108,7 @@ def build(self, input_shape):
 
   def __call__(self, inputs, *args, **kwargs):
     if not isinstance(inputs, random_variable.RandomVariable):
-      return super(DiscreteAutoregressiveFlow, self).__call__(  # pytype: disable=attribute-error  # typed-keras
+      return super(DiscreteAutoregressiveFlow, self).__call__(
           inputs, *args, **kwargs)
     return transformed_random_variable.TransformedRandomVariable(inputs, self)
 
@@ -310,7 +310,7 @@ def build(self, input_shape):
 
   def __call__(self, inputs, *args, **kwargs):
     if not isinstance(inputs, random_variable.RandomVariable):
-      return super(DiscreteBipartiteFlow, self).__call__(  # pytype: disable=attribute-error  # typed-keras
+      return super(DiscreteBipartiteFlow, self).__call__(
           inputs, *args, **kwargs)
     return transformed_random_variable.TransformedRandomVariable(inputs, self)
 
@@ -404,7 +404,7 @@ def build(self, input_shape):
 
   def __call__(self, inputs, *args, **kwargs):
     if not isinstance(inputs, random_variable.RandomVariable):
-      return super(SinkhornAutoregressiveFlow, self).__call__(  # pytype: disable=attribute-error  # typed-keras
+      return super(SinkhornAutoregressiveFlow, self).__call__(
           inputs, *args, **kwargs)
     return transformed_random_variable.TransformedRandomVariable(inputs, self)
 
diff --git a/edward2/tensorflow/layers/embeddings.py b/edward2/tensorflow/layers/embeddings.py
index 4ebdc07d..945d2745 100644
--- a/edward2/tensorflow/layers/embeddings.py
+++ b/edward2/tensorflow/layers/embeddings.py
@@ -92,7 +92,7 @@ def __init__(self,
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.embeddings_initializer, tf.keras.layers.Layer):
-      self.embeddings = self.embeddings_initializer(self.embeddings.shape,  # pyrefly: ignore[not-callable]
+      self.embeddings = self.embeddings_initializer(self.embeddings.shape,
                                                     self.dtype)
 
   def call(self, *args, **kwargs):
diff --git a/edward2/tensorflow/layers/gaussian_process.py b/edward2/tensorflow/layers/gaussian_process.py
index c6f5626b..fd4c674b 100644
--- a/edward2/tensorflow/layers/gaussian_process.py
+++ b/edward2/tensorflow/layers/gaussian_process.py
@@ -407,11 +407,11 @@ def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.inducing_inputs_initializer, tf.keras.layers.Layer):
       assert self.conditional_inputs is not None
-      self.conditional_inputs = self.inducing_inputs_initializer(  # pyrefly: ignore[not-callable]
+      self.conditional_inputs = self.inducing_inputs_initializer(
           self.conditional_inputs.shape, self.dtype)
     if isinstance(self.inducing_outputs_initializer, tf.keras.layers.Layer):
       assert self.conditional_outputs is not None
-      self.conditional_outputs = self.inducing_outputs_initializer(  # pyrefly: ignore[not-callable]
+      self.conditional_outputs = self.inducing_outputs_initializer(
           self.conditional_outputs.shape, self.dtype)
 
   def call(self, inputs):
diff --git a/edward2/tensorflow/layers/heteroscedastic.py b/edward2/tensorflow/layers/heteroscedastic.py
index 69b819c2..7af5c1c3 100644
--- a/edward2/tensorflow/layers/heteroscedastic.py
+++ b/edward2/tensorflow/layers/heteroscedastic.py
@@ -303,7 +303,7 @@ def call(self, inputs, training=True, seed=None):
       if seed is not None:
         raise ValueError('Seed should not be provided when running in graph '
                          'mode, but %s was provided.' % seed)
-    with tf.name_scope(self._name):  # pyrefly: ignore[bad-instantiation]
+    with tf.name_scope(self._name):
       locs = self._compute_loc_param(inputs)  # pylint: disable=assignment-from-none
       scale = self._compute_scale_param(inputs)  # pylint: disable=assignment-from-none
 
@@ -482,7 +482,7 @@ def _compute_loc_param(self, inputs):
     Returns:
       Tensor of shape [batch_size, num_classes].
     """
-    return self._loc_layer(inputs)  # pyrefly: ignore[not-callable]
+    return self._loc_layer(inputs)
 
   def _compute_scale_param(self, inputs):
     """Computes scale parameter of the "logits distribution".
@@ -493,7 +493,7 @@ def _compute_scale_param(self, inputs):
     Returns:
       Tensor of shape [batch_size, num_classes].
     """
-    return self._scale_layer(inputs) + MIN_SCALE_MONTE_CARLO  # pyrefly: ignore[not-callable]
+    return self._scale_layer(inputs) + MIN_SCALE_MONTE_CARLO
 
   def get_config(self):
     config = {
@@ -655,7 +655,7 @@ def _compute_loc_param(self, inputs):
     Returns:
       Tensor of shape [batch_size, num_classes].
     """
-    return self._loc_layer(inputs)  # pyrefly: ignore[not-callable]
+    return self._loc_layer(inputs)
 
   def _compute_scale_param(self, inputs):
     """Computes scale parameter of the "logits distribution".
@@ -668,10 +668,10 @@ def _compute_scale_param(self, inputs):
       [batch_size, num_classes]).
     """
     if self._parameter_efficient:
-      return (inputs, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+      return (inputs, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
     else:
-      return (self._scale_layer(inputs),  # pyrefly: ignore[not-callable]
-              self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+      return (self._scale_layer(inputs),
+              self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
 
   def _compute_diagonal_noise_samples(self, diag_scale, num_samples, seed):
     """Compute samples of the diagonal elements logit noise.
@@ -772,9 +772,9 @@ def _compute_noise_samples(self, scale, num_samples, seed):
         factor_loadings, num_samples, seed)
 
     if self._parameter_efficient:
-      res = self._scale_layer_homoscedastic(standard_normal_samples)  # pyrefly: ignore[not-callable]
+      res = self._scale_layer_homoscedastic(standard_normal_samples)
       res *= tf.expand_dims(
-          self._scale_layer_heteroscedastic(factor_loadings), 1)  # pyrefly: ignore[not-callable]
+          self._scale_layer_heteroscedastic(factor_loadings), 1)
     else:
       # reshape scale vector into factor loadings matrix
       factor_loadings = tf.reshape(factor_loadings,
@@ -957,7 +957,7 @@ def _compute_mc_samples(self, locs, scale, num_samples, seed):
     locs = tf.expand_dims(locs, axis=1)
     noise_samples = self._compute_noise_samples(scale, num_samples, seed)
     latents = locs + noise_samples
-    latents = tf.keras.layers.Reshape(  # pyrefly: ignore[not-callable]
+    latents = tf.keras.layers.Reshape(
         [num_samples, self._ensemble_size, self._num_classes])(latents)
     if self._num_classes == 2:
       return tf.math.sigmoid(latents / self._temperature)
@@ -973,7 +973,7 @@ def _compute_loc_param(self, inputs):
     Returns:
       Tensor of shape [batch_size, num_classes].
     """
-    return self._loc_layer(inputs)  # pyrefly: ignore[not-callable]
+    return self._loc_layer(inputs)
 
   def _compute_scale_param(self, inputs):
     """Computes scale parameter of the "logits distribution".
@@ -986,10 +986,10 @@ def _compute_scale_param(self, inputs):
       [batch_size, num_classes]).
     """
     if self._parameter_efficient:
-      return (inputs, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+      return (inputs, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
     else:
-      return (self._scale_layer(inputs),  # pyrefly: ignore[not-callable]
-              self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+      return (self._scale_layer(inputs),
+              self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
 
   def _compute_diagonal_noise_samples(self, diag_scale, num_samples, seed):
     """Compute samples of the diagonal elements logit noise.
@@ -1092,9 +1092,9 @@ def _compute_noise_samples(self, scale, num_samples, seed):
         factor_loadings, num_samples, seed)
 
     if self._parameter_efficient:
-      res = self._scale_layer_homoscedastic(standard_normal_samples)  # pyrefly: ignore[not-callable]
+      res = self._scale_layer_homoscedastic(standard_normal_samples)
       res *= tf.expand_dims(
-          self._scale_layer_heteroscedastic(factor_loadings), 1)  # pyrefly: ignore[not-callable]
+          self._scale_layer_heteroscedastic(factor_loadings), 1)
     else:
       # reshape scale vector into factor loadings matrix
       factor_loadings = tf.reshape(
@@ -1532,7 +1532,7 @@ def _compute_loc_param(self, inputs):
     Returns:
       Tensor of shape [batch_size, num_outputs].
     """
-    return self._loc_layer(inputs)  # pyrefly: ignore[not-callable]
+    return self._loc_layer(inputs)
 
   def _compute_scale_param(self, inputs):
     """Computes scale parameter of the "logits distribution".
@@ -1546,12 +1546,12 @@ def _compute_scale_param(self, inputs):
     """
     if self._num_factors > 0:
       if self._parameter_efficient:
-        return (inputs, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+        return (inputs, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
       else:
-        return (self._scale_layer(inputs),  # pyrefly: ignore[not-callable]
-                self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+        return (self._scale_layer(inputs),
+                self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
     else:
-      return (None, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)  # pyrefly: ignore[not-callable]
+      return (None, self._diag_layer(inputs) + MIN_SCALE_MONTE_CARLO)
 
   def _compute_diagonal_noise_samples(self, diag_scale, num_samples, seed):
     """Compute samples of the diagonal elements logit noise.
@@ -1653,9 +1653,9 @@ def _compute_noise_samples(self, scale, num_samples, seed):
         factor_loadings, num_samples, seed)
 
     if self._parameter_efficient:
-      noise_samples = self._scale_layer_homoscedastic(standard_normal_samples)  # pyrefly: ignore[not-callable]
+      noise_samples = self._scale_layer_homoscedastic(standard_normal_samples)
       noise_samples *= tf.expand_dims(
-          self._scale_layer_heteroscedastic(factor_loadings), 1)  # pyrefly: ignore[not-callable]
+          self._scale_layer_heteroscedastic(factor_loadings), 1)
       return noise_samples + diag_noise_samples
     else:
       # reshape scale vector into factor loadings matrix
@@ -1769,9 +1769,9 @@ def __call__(self, inputs, training=True):
       Tuple of (logits, log_probs, probs). Logits can be used with the
       tf.nn.sigmoid_cross_entropy_with_logits loss function.
     """
-    with tf.name_scope(self._name):  # pyrefly: ignore[bad-instantiation]
-      loc = self._loc_layer(inputs)  # pyrefly: ignore[not-callable]
-      scale = self._diag_layer(inputs) + self._min_scale  # pyrefly: ignore[not-callable]
+    with tf.name_scope(self._name):
+      loc = self._loc_layer(inputs)
+      scale = self._diag_layer(inputs) + self._min_scale
       loc = tf.cast(loc, tf.float32)
       scale = tf.cast(scale, tf.float32)
 
@@ -1875,7 +1875,7 @@ def __call__(self, inputs, training=True):
       tf.nn.sigmoid_cross_entropy_with_logits or
       tf.nn.softmax_cross_entropy_with_logits loss functions.
     """
-    with tf.name_scope(self._name):  # pyrefly: ignore[bad-instantiation]
+    with tf.name_scope(self._name):
       if self._averaging == 'ensemble_cross_ent':
         log_weights = tf.math.log(self._ensemble_weighting)
 
@@ -2131,9 +2131,9 @@ def _compute_noise_samples(self, scale: Collection[tf.Tensor],
         factor_loadings, num_samples, seed)
 
     if self._parameter_efficient:
-      res = self._scale_layer_homoscedastic(standard_normal_samples)  # pyrefly: ignore[not-callable]
+      res = self._scale_layer_homoscedastic(standard_normal_samples)
       res *= tf.expand_dims(
-          self._scale_layer_heteroscedastic(factor_loadings), 1)  # pyrefly: ignore[not-callable]
+          self._scale_layer_heteroscedastic(factor_loadings), 1)
     else:
       # reshape scale vector into factor loadings matrix
       batch_size, height, width = tf.shape(factor_loadings)[:3]
diff --git a/edward2/tensorflow/layers/hetsngp.py b/edward2/tensorflow/layers/hetsngp.py
index c7aff0e2..8a5eae71 100644
--- a/edward2/tensorflow/layers/hetsngp.py
+++ b/edward2/tensorflow/layers/hetsngp.py
@@ -165,11 +165,11 @@ def __init__(
     self.sngp_var_weight = sngp_var_weight
     self.het_var_weight = het_var_weight
 
-  def _compute_loc_param(self, inputs, training):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def _compute_loc_param(self, inputs, training):  # pyrefly: ignore[bad-override]
     """Computes the mean logits as the mean-field logits of the SNGP."""
     return self.sngp_layer(inputs)
 
-  def _compute_scale_param(self, inputs, covmat_sngp, training):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def _compute_scale_param(self, inputs, covmat_sngp, training):  # pyrefly: ignore[bad-override]
     """Computes the variances for the logits."""
     low_rank, diag = super()._compute_scale_param(inputs)
     sngp_marginal_vars = tf.expand_dims(tf.linalg.diag_part(covmat_sngp), -1)
@@ -209,7 +209,7 @@ def __call__(self, inputs, training=True, seed=None):
       if seed is not None:
         raise ValueError('Seed should not be provided when running in graph '
                          'mode, but %s was provided.' % seed)
-    with tf.name_scope(self._name):  # pyrefly: ignore[bad-instantiation]
+    with tf.name_scope(self._name):
       locs, covmat_sngp = self._compute_loc_param(inputs, training)  # pylint: disable=assignment-from-none
       scale = self._compute_scale_param(inputs, covmat_sngp, training)  # pylint: disable=assignment-from-none
 
diff --git a/edward2/tensorflow/layers/made.py b/edward2/tensorflow/layers/made.py
index 0f11d782..b17ca238 100644
--- a/edward2/tensorflow/layers/made.py
+++ b/edward2/tensorflow/layers/made.py
@@ -125,7 +125,7 @@ def build(self, input_shape):
     self.network.add(tf.keras.layers.Reshape([length, self.units]))
     self.built = True
 
-  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def call(self, inputs):  # pyrefly: ignore[bad-override]
     return self.network(inputs)
 
 
@@ -219,7 +219,7 @@ def create_masks(input_dim,
 def make_masked_initializer(mask):
   initializer = tf.keras.initializers.GlorotUniform()
   def masked_initializer(shape, dtype=None):
-    return mask * initializer(shape, dtype)  # pyrefly: ignore[not-callable]
+    return mask * initializer(shape, dtype)
   return masked_initializer
 
 
diff --git a/edward2/tensorflow/layers/neural_process.py b/edward2/tensorflow/layers/neural_process.py
index e9f7a7d1..bb5e540f 100644
--- a/edward2/tensorflow/layers/neural_process.py
+++ b/edward2/tensorflow/layers/neural_process.py
@@ -36,9 +36,9 @@ def batch_mlp(inputs, hidden_sizes):
   hidden = tf.reshape(inputs, (-1, filter_size))
 
   for size in hidden_sizes[:-1]:
-    hidden = tf.keras.layers.Dense(size, activation=tf.nn.relu)(hidden)  # pyrefly: ignore[not-callable]
+    hidden = tf.keras.layers.Dense(size, activation=tf.nn.relu)(hidden)
 
-  output = tf.keras.layers.Dense(hidden_sizes[-1], activation=None)(hidden)  # pyrefly: ignore[not-callable]
+  output = tf.keras.layers.Dense(hidden_sizes[-1], activation=None)(hidden)
   output = tf.reshape(output, (batch_size, -1, hidden_sizes[-1]))
   return output
 
@@ -264,11 +264,11 @@ def latent_encoder(self, x, y):
     per_example_embedding = batch_mlp(
         encoder_input, self._latent_encoder_sizes)
     dataset_embedding = tf.reduce_mean(per_example_embedding, axis=1)
-    hidden = tf.keras.layers.Dense(  # pyrefly: ignore[not-callable]
+    hidden = tf.keras.layers.Dense(
         (self._latent_encoder_sizes[-1] + self._num_latents)//2,
         activation=tf.nn.relu)(dataset_embedding)
-    loc = tf.keras.layers.Dense(self._num_latents, activation=None)(hidden)  # pyrefly: ignore[not-callable]
-    untransformed_scale = tf.keras.layers.Dense(self._num_latents,  # pyrefly: ignore[not-callable]
+    loc = tf.keras.layers.Dense(self._num_latents, activation=None)(hidden)
+    untransformed_scale = tf.keras.layers.Dense(self._num_latents,
                                                 activation=None)(hidden)
     # Constraint scale following Garnelo et al. (2018).
     scale_diag = 0.1 + 0.9 * tf.sigmoid(untransformed_scale)
diff --git a/edward2/tensorflow/layers/normalization.py b/edward2/tensorflow/layers/normalization.py
index 7370cecc..8cde0c1f 100644
--- a/edward2/tensorflow/layers/normalization.py
+++ b/edward2/tensorflow/layers/normalization.py
@@ -78,7 +78,7 @@ def __call__(self, inputs, *args, **kwargs):
           1. / (tf.sqrt(variance) + self.epsilon))
 
     if not isinstance(inputs, random_variable.RandomVariable):
-      return super(ActNorm, self).__call__(inputs, *args, **kwargs)  # pytype: disable=attribute-error  # typed-keras
+      return super(ActNorm, self).__call__(inputs, *args, **kwargs)
     return transformed_random_variable.TransformedRandomVariable(inputs, self)
 
   def call(self, inputs):
@@ -325,7 +325,7 @@ def __init__(self,
     super(SpectralNormalization, self).__init__(
         layer, name=wrapper_name, **kwargs)
 
-  def build(self, input_shape):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def build(self, input_shape):  # pyrefly: ignore[bad-override]
     super(SpectralNormalization, self).build(input_shape)
     self.layer.kernel._aggregation = self.aggregation  # pylint: disable=protected-access
     self._dtype = self.layer.kernel.dtype
@@ -440,7 +440,7 @@ def __init__(self,
           .format(input=layer))
     super(SpectralNormalizationConv2D, self).__init__(layer, **kwargs)
 
-  def build(self, input_shape):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def build(self, input_shape):  # pyrefly: ignore[bad-override]
     self.layer.build(input_shape)
     self.layer.kernel._aggregation = self.aggregation  # pylint: disable=protected-access
     self._dtype = self.layer.kernel.dtype
diff --git a/edward2/tensorflow/layers/posterior_network.py b/edward2/tensorflow/layers/posterior_network.py
index 25a4de6c..168bfd44 100644
--- a/edward2/tensorflow/layers/posterior_network.py
+++ b/edward2/tensorflow/layers/posterior_network.py
@@ -54,7 +54,7 @@ def __init__(self,
 
   def build(self, input_shape):
     """Builds the layer based on the passed input shape."""
-    with tf.name_scope(self.name):  # pyrefly: ignore[bad-instantiation]
+    with tf.name_scope(self.name):
       # Using the PyTorch default hyperparameters.
       self.batch_norm = tf.keras.layers.BatchNormalization(epsilon=1e-5,
                                                            momentum=0.9)
@@ -133,7 +133,7 @@ def __init__(self,
         validate_args=validate_args,
         inverse_min_event_ndims=1,
         name=name)
-    with tf.name_scope(name) as name:  # pyrefly: ignore[bad-instantiation]
+    with tf.name_scope(name) as name:
       self._name = name
       if x0 is None:
         x0 = tf.zeros(dim)
diff --git a/edward2/tensorflow/layers/recurrent.py b/edward2/tensorflow/layers/recurrent.py
index 98711ee8..457fb401 100644
--- a/edward2/tensorflow/layers/recurrent.py
+++ b/edward2/tensorflow/layers/recurrent.py
@@ -117,7 +117,7 @@ def build(
         def bias_initializer(_, *args, **kwargs):
           return tf.keras.backend.concatenate([
               self.bias_initializer((self.units,), *args, **kwargs),
-              tf.keras.initializers.Ones()((self.units,), *args, **kwargs),  # pyrefly: ignore[not-callable]
+              tf.keras.initializers.Ones()((self.units,), *args, **kwargs),
               self.bias_initializer((self.units * 2,), *args, **kwargs),
           ])
       else:
@@ -146,14 +146,14 @@ def call(self, *args, **kwargs):
   def call_weights(self):
     """Calls any weights if the initializer is itself a layer."""
     if isinstance(self.kernel_initializer, tf.keras.layers.Layer):
-      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)  # pyrefly: ignore[not-callable]
+      self.kernel = self.kernel_initializer(self.kernel.shape, self.dtype)
     if isinstance(self.recurrent_initializer, tf.keras.layers.Layer):
-      self.recurrent_kernel = self.recurrent_initializer(  # pyrefly: ignore[not-callable]
+      self.recurrent_kernel = self.recurrent_initializer(
           self.recurrent_kernel.shape, self.dtype)
     if isinstance(self.bias_initializer, tf.keras.layers.Layer):
       if self.bias is None:
         raise ValueError('self.bias is None.')
-      self.bias = self.bias_initializer(self.bias.shape, self.dtype)  # pyrefly: ignore[not-callable]
+      self.bias = self.bias_initializer(self.bias.shape, self.dtype)
     self.called_weights = True
 
   def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
@@ -594,7 +594,7 @@ def bias_initializer(_, *args, **kwargs):
           return tf.concat([
               self.bias_initializer([self.ensemble_size, self.units], *args,
                                     **kwargs),
-              tf.keras.initializers.Ones()([self.ensemble_size, self.units],  # pyrefly: ignore[not-callable]
+              tf.keras.initializers.Ones()([self.ensemble_size, self.units],
                                            *args, **kwargs),
               self.bias_initializer([self.ensemble_size, self.units * 2],
                                     *args, **kwargs),
@@ -647,7 +647,7 @@ def _sample_weights(self, inputs=None, batch_size=None, dtype=None):
     # Sample parameters for each input example.
     def sample(weight_variable, weight_initializer, shape):
       if isinstance(weight_initializer, tf.keras.layers.Layer):
-        weights = weight_initializer(  # pyrefly: ignore[not-callable]
+        weights = weight_initializer(
             shape, self.dtype).distribution.sample(examples_per_model)
         weights = tf.transpose(weights, [1, 0, 2])
       else:
diff --git a/edward2/tensorflow/layers/stochastic_output.py b/edward2/tensorflow/layers/stochastic_output.py
index 47392183..4e3b29f9 100644
--- a/edward2/tensorflow/layers/stochastic_output.py
+++ b/edward2/tensorflow/layers/stochastic_output.py
@@ -81,7 +81,7 @@ def build(self, input_shape=None):
     self.built = True
 
   def call(self, inputs):
-    net = self.layer(inputs)  # pyrefly: ignore[not-callable]
+    net = self.layer(inputs)
     logits, loc, scale = tf.split(net, 3, axis=-1)
     if self.logits_constraint:
       logits = self.logits_constraint(logits)
diff --git a/edward2/tensorflow/layers/utils.py b/edward2/tensorflow/layers/utils.py
index e6c840c4..b6f1af2f 100644
--- a/edward2/tensorflow/layers/utils.py
+++ b/edward2/tensorflow/layers/utils.py
@@ -56,13 +56,13 @@ def _add_weight(self,
       if not regularizer.built:
         regularizer.build(shape)
     if isinstance(initializer, tf.keras.layers.Layer):
-      with tf.name_scope(name):  # pyrefly: ignore[bad-instantiation]
-        weight = initializer(shape, dtype)  # pyrefly: ignore[not-callable]
+      with tf.name_scope(name):
+        weight = initializer(shape, dtype)
       if regularizer is not None:
         def loss_fn():
           """Creates a regularization loss `Tensor`."""
-          with tf.name_scope(name + '/Regularizer'):  # pyrefly: ignore[bad-instantiation, unsupported-operation]
-            return regularizer(initializer(shape, dtype))  # pyrefly: ignore[not-callable]
+          with tf.name_scope(name + '/Regularizer'):  # pyrefly: ignore[unsupported-operation]
+            return regularizer(initializer(shape, dtype))
         self.add_loss(loss_fn)
       return weight
     return super(cls, self).add_weight(name=name,
diff --git a/edward2/tensorflow/random_variable.py b/edward2/tensorflow/random_variable.py
index b96028c6..fe85473e 100644
--- a/edward2/tensorflow/random_variable.py
+++ b/edward2/tensorflow/random_variable.py
@@ -148,7 +148,7 @@ def sample_shape_tensor(self, name="sample_shape_tensor"):
     Returns:
       sample_shape: `Tensor`.
     """
-    with tf.name_scope(name):  # pyrefly: ignore[bad-instantiation]
+    with tf.name_scope(name):
       if tf.is_tensor(self._sample_shape):
         return self._sample_shape
       return tf.convert_to_tensor(self.sample_shape.as_list(), dtype=tf.int32)
diff --git a/examples/no_u_turn_sampler/logistic_regression.py b/examples/no_u_turn_sampler/logistic_regression.py
index 0bf2aedb..e6631bc1 100644
--- a/examples/no_u_turn_sampler/logistic_regression.py
+++ b/examples/no_u_turn_sampler/logistic_regression.py
@@ -75,8 +75,8 @@ def logistic_regression(features):
 def covertype():
   """Builds the Covertype data set."""
   data = sklearn.datasets.fetch_covtype()
-  features = data.data  # pytype: disable=attribute-error
-  labels = data.target  # pytype: disable=attribute-error
+  features = data.data
+  labels = data.target
 
   # Normalize features and append a column of ones for the intercept.
   features -= features.mean(0)
diff --git a/examples/no_u_turn_sampler/nuts.py b/examples/no_u_turn_sampler/nuts.py
index f0bd45f1..566f01ec 100644
--- a/examples/no_u_turn_sampler/nuts.py
+++ b/examples/no_u_turn_sampler/nuts.py
@@ -73,8 +73,8 @@ def kernel(target_log_prob_fn,
   if not tf.executing_eagerly():
     raise NotImplementedError("`kernel` is only available in Eager mode.")
 
-  with tf.name_scope(name):  # pyrefly: ignore[bad-instantiation]
-    with tf.name_scope("initialize"):  # pyrefly: ignore[bad-instantiation]
+  with tf.name_scope(name):
+    with tf.name_scope("initialize"):
       current_state = [tf.convert_to_tensor(s) for s in current_state]
       step_size = [tf.convert_to_tensor(s) for s in step_size]
       if (current_target_log_prob is None or
@@ -487,7 +487,7 @@ def _random_bernoulli(shape,
                       seed=None,
                       name="random_bernoulli"):
   """Returns samples from a Bernoulli distribution."""
-  with tf.name_scope(name):  # pyrefly: ignore[bad-instantiation]
+  with tf.name_scope(name):
     probs = tf.convert_to_tensor(probs)
     random_uniform = tf.random.uniform(shape, dtype=probs.dtype, seed=seed)
     return tf.cast(tf.less(random_uniform, probs), dtype)