Replace # pytype: disable suppressions with # pyrefly: ignore

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(

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)

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.

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})

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])

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]

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)
 

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):

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):