Internal change

edward2/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/attention.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -182,12 +182,12 @@ def f(x, *args, **kwargs):
         query,
         key,
         value,
-        mask=mask,
-        dropout_rng=dropout_rng,
-        dropout_rate=self.dropout_rate,
-        broadcast_dropout=self.broadcast_dropout,
-        deterministic=deterministic,
-        dtype=self.dtype,
+        mask=mask,  # pyrefly: ignore[unexpected-keyword]
+        dropout_rng=dropout_rng,  # pyrefly: ignore[unexpected-keyword]
+        dropout_rate=self.dropout_rate,  # pyrefly: ignore[unexpected-keyword]
+        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
     # back to the original inputs dimensions
     def dense_fn2(name):

edward2/jax/nn/attention_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/dense.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/dense_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/heteroscedastic_lib.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -257,7 +257,7 @@ def _compute_noise_samples(self, scale, num_samples):
   def get_temperature(self):
     if self.tune_temperature:
       return compute_temperature(
-          self._pre_sigmoid_temperature,
+          self._pre_sigmoid_temperature,  # pyrefly: ignore[bad-argument-type]
           lower=self.temperature_lower_bound,
           upper=self.temperature_upper_bound)
     else:
@@ -510,7 +510,7 @@ def _compute_noise_samples(self, scale, num_samples):
   def get_temperature(self):
     if self.tune_temperature:
       return compute_temperature(
-          self._pre_sigmoid_temperature,
+          self._pre_sigmoid_temperature,  # pyrefly: ignore[bad-argument-type]
           lower=self.temperature_lower_bound,
           upper=self.temperature_upper_bound)
     else:
@@ -609,33 +609,33 @@ def setup(self):
       self.actual_latent_dim = self.latent_dim
 
     if self.parameter_efficient:
-      self._scale_layer_homoscedastic = dense.DenseBatchEnsemble(
+      self._scale_layer_homoscedastic = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.actual_latent_dim,
           ens_size=self.ens_size,
           alpha_init=self.alpha_init,
           gamma_init=self.gamma_init,
           kernel_init=self.kernel_init,
           name='scale_layer_homoscedastic')
-      self._scale_layer_heteroscedastic = dense.DenseBatchEnsemble(
+      self._scale_layer_heteroscedastic = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.actual_latent_dim,
           ens_size=self.ens_size,
           alpha_init=self.alpha_init,
           gamma_init=self.gamma_init,
           kernel_init=self.kernel_init,
           name='scale_layer_heteroscedastic')
     elif self.num_factors > 0:
-      self._scale_layer = dense.DenseBatchEnsemble(
+      self._scale_layer = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.actual_latent_dim * self.num_factors,
           ens_size=self.ens_size,
           alpha_init=self.alpha_init,
           gamma_init=self.gamma_init,
           kernel_init=self.kernel_init,
           name='scale_layer')
 
-    self._loc_layer = dense.DenseBatchEnsemble(self.num_classes,
+    self._loc_layer = dense.DenseBatchEnsemble(self.num_classes,  # pyrefly: ignore[bad-assignment]
                                                ens_size=self.ens_size,
                                                name='loc_layer')
-    self._diag_layer = dense.DenseBatchEnsemble(self.actual_latent_dim,
+    self._diag_layer = dense.DenseBatchEnsemble(self.actual_latent_dim,  # pyrefly: ignore[bad-assignment]
                                                 ens_size=self.ens_size,
                                                 name='diag_layer')
 
@@ -664,36 +664,36 @@ def setup(self):
       self.actual_latent_dim = self.latent_dim
 
     if self.parameter_efficient:
-      self._scale_layer_homoscedastic = dense.DenseBatchEnsemble(
+      self._scale_layer_homoscedastic = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.actual_latent_dim,
           ens_size=self.ens_size,
           alpha_init=self.alpha_init,
           gamma_init=self.gamma_init,
           kernel_init=self.kernel_init,
           name='scale_layer_homoscedastic')
-      self._scale_layer_heteroscedastic = dense.DenseBatchEnsemble(
+      self._scale_layer_heteroscedastic = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.actual_latent_dim,
           ens_size=self.ens_size,
           alpha_init=self.alpha_init,
           gamma_init=self.gamma_init,
           kernel_init=self.kernel_init,
           name='scale_layer_heteroscedastic')
     elif self.num_factors > 0:
-      self._scale_layer = dense.DenseBatchEnsemble(
+      self._scale_layer = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.actual_latent_dim * self.num_factors,
           ens_size=self.ens_size,
           alpha_init=self.alpha_init,
           gamma_init=self.gamma_init,
           kernel_init=self.kernel_init,
           name='scale_layer')
 
-    self._loc_layer = dense.DenseBatchEnsemble(self.num_outputs,
+    self._loc_layer = dense.DenseBatchEnsemble(self.num_outputs,  # pyrefly: ignore[bad-assignment]
                                                ens_size=self.ens_size,
                                                alpha_init=self.alpha_init,
                                                gamma_init=self.gamma_init,
                                                kernel_init=self.kernel_init,
                                                name='loc_layer')
-    self._diag_layer = dense.DenseBatchEnsemble(self.actual_latent_dim,
+    self._diag_layer = dense.DenseBatchEnsemble(self.actual_latent_dim,  # pyrefly: ignore[bad-assignment]
                                                 ens_size=self.ens_size,
                                                 alpha_init=self.alpha_init,
                                                 gamma_init=self.gamma_init,

edward2/jax/nn/normalization.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -187,7 +187,7 @@ def scan_body(carry, _):
 
 class SpectralNormalizationConv2D(SpectralNormalization):
   __doc__ = "Implements spectral normalization for Conv layers based on [2].\n" + "\n".join(
-      SpectralNormalization.__doc__.split("\n")[1:])
+      SpectralNormalization.__doc__.split("\n")[1:])  # pyrefly: ignore[missing-attribute]
 
   kernel_apply_kwargs: Mapping[str, Any] = flax.core.FrozenDict(
       feature_group_count=1, padding="SAME", use_bias=False)
@@ -208,8 +208,8 @@ def _abs_sq(x):
   # promote x to at least float32, this avoids half precision computation
   # but preserves double or complex floating points
   x = jnp.asarray(x, jnp.promote_types(jnp.float32, jnp.result_type(x)))
-  mean = jnp.mean(x, axes)
-  mean2 = jnp.mean(_abs_sq(x), axes)
+  mean = jnp.mean(x, axes)  # pyrefly: ignore[bad-argument-type]
+  mean2 = jnp.mean(_abs_sq(x), axes)  # pyrefly: ignore[bad-argument-type]
   if axis_name is not None:
     concatenated_mean = jnp.concatenate([mean, mean2])
     mean, mean2 = jnp.split(

edward2/jax/nn/normalization_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/random_feature.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -211,7 +211,7 @@ def __call__(self, inputs: Array) -> Array:
                               (contracting_dims, batch_dims))
     outputs = outputs + jnp.broadcast_to(bias.value, outputs.shape)
 
-    return self._feature_scale * self.activation(outputs)
+    return self._feature_scale * self.activation(outputs)  # pyrefly: ignore[unsupported-operation]
 
 
 class LaplaceRandomFeatureCovariance(nn.Module):
@@ -340,10 +340,10 @@ def update_precision_matrix(self, gp_features: Array,
 
     # Computes precision matrix within new batch.
     if self.likelihood == 'binary_logistic':
-      prob = nn.sigmoid(gp_logits)
+      prob = nn.sigmoid(gp_logits)  # pyrefly: ignore[bad-argument-type]
       prob_multiplier = prob * (1. - prob)
     elif self.likelihood == 'poisson':
-      prob_multiplier = jnp.exp(gp_logits)
+      prob_multiplier = jnp.exp(gp_logits)  # pyrefly: ignore[bad-argument-type]
     else:
       prob_multiplier = 1.
 
@@ -703,16 +703,16 @@ class MCSigmoidDenseFASNGPBE(MCSigmoidDenseFASNGP):
 
   def setup(self):
     if self.parameter_efficient:
-      self._scale_layer_homoscedastic = dense.DenseBatchEnsemble(
+      self._scale_layer_homoscedastic = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.num_outputs,
           ens_size=self.ens_size,
           name='scale_layer_homoscedastic')
-      self._scale_layer_heteroscedastic = dense.DenseBatchEnsemble(
+      self._scale_layer_heteroscedastic = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.num_outputs,
           ens_size=self.ens_size,
           name='scale_layer_heteroscedastic')
     elif self.num_factors > 0:
-      self._scale_layer = dense.DenseBatchEnsemble(
+      self._scale_layer = dense.DenseBatchEnsemble(  # pyrefly: ignore[bad-assignment]
           self.num_outputs * self.num_factors,
           ens_size=self.ens_size,
           name='scale_layer')
@@ -726,6 +726,6 @@ def setup(self):
         output_kwargs=self.output_kwargs,
         covmat_kwargs=self.covmat_kwargs,
         name='loc_layer')
-    self._diag_layer = dense.DenseBatchEnsemble(self.num_outputs,
+    self._diag_layer = dense.DenseBatchEnsemble(self.num_outputs,  # pyrefly: ignore[bad-assignment]
                                                 ens_size=self.ens_size,
                                                 name='diag_layer')

edward2/jax/nn/random_feature_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/utils.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/jax/nn/utils_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/maps.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -29,7 +29,7 @@
 
 
 @overload
-def robust_map(
+def robust_map(  # pyrefly: ignore[inconsistent-overload]
     fn: Callable[[T], U],
     inputs: Sequence[T],
     error_output: V = ...,
@@ -45,7 +45,7 @@ def robust_map(
 
 
 @overload
-def robust_map(
+def robust_map(  # pyrefly: ignore[inconsistent-overload]
     fn: Callable[[T], U],
     inputs: Sequence[T],
     error_output: V = ...,
@@ -61,7 +61,7 @@ def robust_map(
 
 
 @overload
-def robust_map(
+def robust_map(  # pyrefly: ignore[inconsistent-overload]
     fn: Callable[[T], U],
     inputs: Sequence[T],
     error_output: V = ...,
@@ -77,7 +77,7 @@ def robust_map(
 
 
 @overload
-def robust_map(
+def robust_map(  # pyrefly: ignore[inconsistent-overload]
     fn: Callable[[T], U],
     inputs: Sequence[T],
     error_output: V = ...,

edward2/maps_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/numpy/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/numpy/generated_random_variables.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -32,7 +32,7 @@
                             scipy.stats.rv_discrete,
                             scipy.stats.rv_histogram)):
     if hasattr(candidate, 'rvs'):
-      candidate.rvs = traceable(candidate.rvs)
+      candidate.rvs = traceable(candidate.rvs)  # pyrefly: ignore[missing-attribute]
       _globals[candidate_name] = candidate
       __all__.append(candidate_name)
 

edward2/numpy/generated_random_variables_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/numpy/program_transformations.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -124,7 +124,7 @@ def tracer(rv_call, *rv_args, **rv_kwargs):
       rv_kwargs.pop("size", None)
       rv_kwargs.pop("random_state", None)
       rv_kwargs.pop("name", None)
-      log_prob = np.sum(log_prob_fn(cls(), value, *rv_args, **rv_kwargs))
+      log_prob = np.sum(log_prob_fn(cls(), value, *rv_args, **rv_kwargs))  # pyrefly: ignore[not-callable]
       log_probs.append(log_prob)
       return value
 

edward2/numpy/program_transformations_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

edward2/tensorflow/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2025 The Edward2 Authors.
+# Copyright 2026 The Edward2 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.