Internal change

PiperOrigin-RevId: 941350466

Author: h-joo

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.