Automated Code Change

sonnet/__init__.py

@@ -160,6 +160,6 @@
 #                 ||     ||
 #
 try:
-  del src  # pylint: disable=undefined-variable
+  del src  # pylint: disable=undefined-variable  # pyrefly: ignore[unbound-name]
 except NameError:
   pass

sonnet/src/base.py

@@ -58,7 +58,7 @@ def no_name_scope(method: T) -> T:
 class ModuleMetaclass(abc.ABCMeta):
   """Metaclass for `Module`."""
 
-  def __new__(
+  def __new__(  # pyrefly: ignore[invalid-annotation]
       cls: Type[Type[T]],
       name: str,
       bases: Tuple[Type[Any], ...],
@@ -89,7 +89,7 @@ def __new__(
 
     clsdict.setdefault("__repr__", lambda module: module._auto_repr)  # pylint: disable=protected-access
 
-    new_cls = super(ModuleMetaclass, cls).__new__(cls, name, bases, clsdict)  # pylint: disable=too-many-function-args
+    new_cls = super(ModuleMetaclass, cls).__new__(cls, name, bases, clsdict)  # pylint: disable=too-many-function-args  # pyrefly: ignore[invalid-argument]
 
     for method_name in methods:
       # Note: the below is quite subtle, we need to ensure that we're wrapping
@@ -127,7 +127,7 @@ def __call__(cls: Type[T], *args, **kwargs) -> T:
       ctor_name_scope = getattr(module, "_ctor_name_scope", None)
       if ctor_name_scope is not None:
         ctor_name_scope.__exit__(*exc_info)
-        del module._ctor_name_scope
+        del module._ctor_name_scope  # pyrefly: ignore[missing-attribute]
 
       # TODO(tomhennigan) Remove `_scope_name` after next TF release.
       ran_super_ctor = (
@@ -139,7 +139,7 @@ def __call__(cls: Type[T], *args, **kwargs) -> T:
             "is not supported. Add the following as the first line in your "
             "__init__ method:\n\nsuper(%s, self).__init__()" % cls.__name__)
 
-    module._auto_repr = auto_repr(cls, *args, **kwargs)  # pylint: disable=protected-access
+    module._auto_repr = auto_repr(cls, *args, **kwargs)  # pylint: disable=protected-access  # pyrefly: ignore[missing-attribute]
 
     return module
 
@@ -357,12 +357,12 @@ def allow_empty_variables(module_or_cls: T) -> T:
 
 
 def assert_tf2():
-  if not assert_tf2.checked:
+  if not assert_tf2.checked:  # pyrefly: ignore[missing-attribute]
     with tf.init_scope():
       assert tf.executing_eagerly(), "Sonnet v2 requires TensorFlow 2"
-    assert_tf2.checked = True
+    assert_tf2.checked = True  # pyrefly: ignore[missing-attribute]
 
-assert_tf2.checked = False
+assert_tf2.checked = False  # pyrefly: ignore[missing-attribute]
 
 
 class Module(tf.Module, metaclass=ModuleMetaclass):

sonnet/src/batch_apply.py

@@ -127,7 +127,7 @@ def split_leading_dim(
 
 def maybe_prod(s: Sequence[Union[int, None]]) -> Optional[int]:
   try:
-    return np.prod(s)
+    return np.prod(s)  # pyrefly: ignore[no-matching-overload]
   except TypeError:
     # Can happen if the input contains `None`.
     return None

sonnet/src/bias.py

@@ -92,13 +92,13 @@ def _initialize(self, inputs):
     utils.assert_minimum_rank(inputs, 2)
 
     input_shape = inputs.shape
-    bias_shape = calculate_bias_shape(input_shape, self.bias_dims)
+    bias_shape = calculate_bias_shape(input_shape, self.bias_dims)  # pyrefly: ignore[bad-argument-type]
 
     input_size = input_shape[1:]
     if self.output_size is not None:
       if self.output_size != input_size:
         raise ValueError("Input shape must be {} not {}".format(
-            (-1,) + self.output_size, input_shape))
+            (-1,) + self.output_size, input_shape))  # pyrefly: ignore[unsupported-operation]
 
     self.input_size = input_size
     self.b = tf.Variable(self.b_init(bias_shape, inputs.dtype), name="b")
@@ -141,10 +141,10 @@ def calculate_bias_shape(input_shape: types.ShapeLike,
     ValueError: If the user attempts to add bias over the mini-batch dimension,
         e.g. `bias_dims=[0]`.
   """
-  input_rank = len(input_shape)
+  input_rank = len(input_shape)  # pyrefly: ignore[bad-argument-type]
   if bias_dims is None:
     # If None, default is to use all dimensions.
-    return input_shape[1:]
+    return input_shape[1:]  # pyrefly: ignore[bad-index]
 
   elif not bias_dims:
     # If empty list, use a scalar bias.
@@ -165,7 +165,7 @@ def calculate_bias_shape(input_shape: types.ShapeLike,
             "Dimension %d (bias_dims=%r) out of range for input of rank %r." %
             (dim, tuple(bias_dims), input_rank))
 
-      bias_shape[dim] = input_shape[dim]
+      bias_shape[dim] = input_shape[dim]  # pyrefly: ignore[bad-index]
     # Strip leading unit dimensions.
     start = input_rank
     for dim in range(1, input_rank):

sonnet/src/conformance/descriptors.py

@@ -55,7 +55,7 @@ def __call__(self, x: tf.Tensor):
       return self.wrapped(x, initial_state)
     else:
       x = tf.expand_dims(x, axis=0)
-      return self.unroller(self.wrapped, x, initial_state)
+      return self.unroller(self.wrapped, x, initial_state)  # pyrefly: ignore[not-callable]
 
 
 def unwrap(module: snt.Module) -> snt.Module:

sonnet/src/conformance/goldens.py

@@ -26,7 +26,7 @@
 
 
 def named_goldens() -> Sequence[Tuple[str, "Golden"]]:
-  return ((name, cls()) for _, name, cls in list_goldens())
+  return ((name, cls()) for _, name, cls in list_goldens())  # pyrefly: ignore[bad-return]
 
 
 def all_goldens(test_method):

sonnet/src/conv.py

@@ -88,7 +88,7 @@ def __init__(self,
       self.padding_func = padding
 
     self.data_format = data_format
-    self._channel_index = utils.get_channel_index(data_format)
+    self._channel_index = utils.get_channel_index(data_format)  # pyrefly: ignore[bad-argument-type]
     self.with_bias = with_bias
 
     self.w_init = w_init

sonnet/src/conv_transpose.py

@@ -110,7 +110,7 @@ def __init__(self,
       raise TypeError("ConvNDTranspose only takes string padding, please "
                       "provide either `SAME` or `VALID`.")
     self._data_format = data_format
-    self._channel_index = utils.get_channel_index(data_format)
+    self._channel_index = utils.get_channel_index(data_format)  # pyrefly: ignore[bad-argument-type]
     self._with_bias = with_bias
 
     self._w_init = w_init
@@ -154,14 +154,14 @@ def _initialize(self, inputs):
     self._dtype = inputs.dtype
 
     if self._output_shape is not None:
-      if len(self._output_shape) != self._num_spatial_dims:
+      if len(self._output_shape) != self._num_spatial_dims:  # pyrefly: ignore[bad-argument-type]
         raise ValueError(
             "The output_shape must be of length {} but instead was {}.".format(
-                self._num_spatial_dims, len(self._output_shape)))
+                self._num_spatial_dims, len(self._output_shape)))  # pyrefly: ignore[bad-argument-type]
       if self._channel_index == 1:
-        self._output_shape = [self._output_channels] + list(self._output_shape)
+        self._output_shape = [self._output_channels] + list(self._output_shape)  # pyrefly: ignore[bad-argument-type]
       else:
-        self._output_shape = list(self._output_shape) + [self._output_channels]
+        self._output_shape = list(self._output_shape) + [self._output_channels]  # pyrefly: ignore[bad-argument-type]
 
     self.w = self._make_w()
     if self._with_bias:

sonnet/src/custom_getter.py

@@ -87,7 +87,7 @@ def _custom_getter(
     orig_getattribute = cls.__getattribute__  # pytype: disable=attribute-error
 
     def new_getattribute(obj, name, orig_getattribute=orig_getattribute):
-      attr = orig_getattribute(obj, name)
+      attr = orig_getattribute(obj, name)  # pyrefly: ignore[bad-argument-count]
 
       if (instances is None) or (obj in instances):
         return getter(attr)

sonnet/src/embed.py

@@ -80,7 +80,7 @@ def __init__(self,
 
     if existing_vocab is None:
       if embed_dim is None:
-        embed_dim = embedding_dim(vocab_size)
+        embed_dim = embedding_dim(vocab_size)  # pyrefly: ignore[bad-argument-type]
       if initializer is None:
         initializer = initializers.TruncatedNormal()
       vocab = initializer([vocab_size, embed_dim], dtype)

sonnet/src/functional/haiku.py

@@ -64,7 +64,7 @@ def notify(f):
   """Wraps `f` such that callbacks are notified about it being called."""
   @functools.wraps(f)
   def wrapper(self, *args, **kwargs):
-    TensorVariableCallbacks.instance.notify(self)
+    TensorVariableCallbacks.instance.notify(self)  # pyrefly: ignore[missing-attribute]
     return f(self, *args, **kwargs)  # pytype: disable=wrong-arg-count
   return wrapper
 
@@ -252,7 +252,7 @@ def getter(next_getter, **kwargs):
 @contextlib.contextmanager
 def track_tensor_variables():
   tensor_variables = []
-  with TensorVariableCallbacks.instance(tensor_variables.append):  # pylint: disable=not-callable
+  with TensorVariableCallbacks.instance(tensor_variables.append):  # pylint: disable=not-callable  # pyrefly: ignore[not-callable]
     yield tensor_variables
 
 
@@ -276,7 +276,7 @@ def callback(v):
     if r not in var_state:
       var_state[r] = (v.initial_tensor_value, v.tensor_value)
 
-  with TensorVariableCallbacks.instance(callback):  # pylint: disable=not-callable
+  with TensorVariableCallbacks.instance(callback):  # pylint: disable=not-callable  # pyrefly: ignore[not-callable]
     yield var_state
 
 

sonnet/src/functional/jax.py

@@ -64,7 +64,7 @@ def wrapper(*args, **kwargs):
       out, aux = out
     grads = tape.gradient(out, params)
     if has_aux:
-      return (out, aux), grads
+      return (out, aux), grads  # pyrefly: ignore[unbound-name]
     else:
       return out, grads
   return wrapper

sonnet/src/initializers.py

@@ -182,12 +182,12 @@ def __call__(self, shape: types.ShapeLike, dtype: tf.DType) -> tf.Tensor:
       raise ValueError("The tensor to initialize must be "
                        "at least two-dimensional")
     elif rank == 2:
-      initializer = tf.eye(num_rows=shape[0], num_columns=shape[1], dtype=dtype)
+      initializer = tf.eye(num_rows=shape[0], num_columns=shape[1], dtype=dtype)  # pyrefly: ignore[bad-index]
     else:  # rank > 2
       initializer = tf.eye(
-          num_rows=shape[-2],
-          num_columns=shape[-1],
-          batch_shape=shape[:-2],
+          num_rows=shape[-2],  # pyrefly: ignore[bad-index]
+          num_columns=shape[-1],  # pyrefly: ignore[bad-index]
+          batch_shape=shape[:-2],  # pyrefly: ignore[bad-index]
           dtype=dtype)
     return self.gain * initializer
 
@@ -223,15 +223,15 @@ def __init__(self, gain: float = 1.0, seed: Optional[int] = None):
 
   def __call__(self, shape: types.ShapeLike, dtype: tf.DType) -> tf.Tensor:
     dtype = _as_floating_dtype(dtype)
-    if len(shape) < 2:
+    if len(shape) < 2:  # pyrefly: ignore[bad-argument-type]
       raise ValueError("The tensor to initialize must be "
                        "at least two-dimensional")
     # Flatten the input shape with the last dimension remaining
     # its original shape so it works for conv2d
     num_rows = 1
-    for dim in shape[:-1]:
+    for dim in shape[:-1]:  # pyrefly: ignore[bad-index]
       num_rows *= dim
-    num_cols = shape[-1]
+    num_cols = shape[-1]  # pyrefly: ignore[bad-index]
     flat_shape = [
         tf.maximum(num_cols, num_rows),
         tf.minimum(num_cols, num_rows)
@@ -370,21 +370,21 @@ def _compute_fans(shape: types.ShapeLike):
   Returns:
     A tuple of scalars `(fan_in, fan_out)`.
   """
-  if len(shape) < 1:  # Just to avoid errors for constants.
+  if len(shape) < 1:  # Just to avoid errors for constants.  # pyrefly: ignore[bad-argument-type]
     fan_in = fan_out = 1
-  elif len(shape) == 1:
-    fan_in = fan_out = shape[0]
-  elif len(shape) == 2:
-    fan_in = shape[0]
-    fan_out = shape[1]
+  elif len(shape) == 1:  # pyrefly: ignore[bad-argument-type]
+    fan_in = fan_out = shape[0]  # pyrefly: ignore[bad-index]
+  elif len(shape) == 2:  # pyrefly: ignore[bad-argument-type]
+    fan_in = shape[0]  # pyrefly: ignore[bad-index]
+    fan_out = shape[1]  # pyrefly: ignore[bad-index]
   else:
     # Assuming convolution kernels (2D, 3D, or more).
     # kernel shape: (..., input_depth, depth)
     receptive_field_size = 1.
-    for dim in shape[:-2]:
+    for dim in shape[:-2]:  # pyrefly: ignore[bad-index]
       receptive_field_size *= dim
-    fan_in = shape[-2] * receptive_field_size
-    fan_out = shape[-1] * receptive_field_size
+    fan_in = shape[-2] * receptive_field_size  # pyrefly: ignore[bad-index]
+    fan_out = shape[-1] * receptive_field_size  # pyrefly: ignore[bad-index]
   return fan_in, fan_out
 
 

sonnet/src/metrics.py

@@ -22,7 +22,7 @@
 import tensorflow as tf
 
 
-class Metric(base.Module, metaclass=abc.ABCMeta):
+class Metric(base.Module, metaclass=abc.ABCMeta):  # pyrefly: ignore[invalid-inheritance]
   """Metric base class."""
 
   @abc.abstractmethod

sonnet/src/nets/dnc/control.py

@@ -83,7 +83,7 @@ def __call__(self, inputs, prev_state):
       output = self._activation(output)
     return output, prev_state
 
-  def initial_state(self, batch_size):
+  def initial_state(self, batch_size):  # pyrefly: ignore[bad-override]
     return tf.zeros([batch_size, 1], dtype=self.dtype)
 
 
@@ -110,6 +110,6 @@ def deep_core(control_name,
       for i in range(num_layers)
   ]
   if skip_connections:
-    return recurrent.deep_rnn_with_skip_connections(cores, name=name)
+    return recurrent.deep_rnn_with_skip_connections(cores, name=name)  # pyrefly: ignore[bad-argument-type]
   else:
     return recurrent.DeepRNN(cores, name=name)

sonnet/src/nets/resnet.py

@@ -41,7 +41,7 @@ def __init__(self,
     self._bn_config = bn_config
 
     batchnorm_args = {"create_scale": True, "create_offset": True}
-    batchnorm_args.update(bn_config)
+    batchnorm_args.update(bn_config)  # pyrefly: ignore[no-matching-overload]
 
     if self._use_projection:
       self._proj_conv = conv.Conv2D(
@@ -120,7 +120,7 @@ def __init__(self,
     self._bn_config = bn_config
 
     batchnorm_args = {"create_scale": True, "create_offset": True}
-    batchnorm_args.update(bn_config)
+    batchnorm_args.update(bn_config)  # pyrefly: ignore[no-matching-overload]
 
     if self._use_projection:
       self._proj_conv = conv.Conv2D(
@@ -274,7 +274,7 @@ def __init__(self,
           create_scale=True,
           create_offset=True,
           name="initial_batchnorm",
-          **bn_config)
+          **bn_config)  # pyrefly: ignore[bad-argument-type]
 
     self._block_groups = []
     strides = [1, 2, 2, 2]
@@ -293,7 +293,7 @@ def __init__(self,
           create_scale=True,
           create_offset=True,
           name="final_batchnorm",
-          **bn_config)
+          **bn_config)  # pyrefly: ignore[bad-argument-type]
 
     self._logits = linear.Linear(
         output_size=num_classes, w_init=initializers.Zeros(), name="logits")

sonnet/src/optimizers/optimizer_tests.py

@@ -37,7 +37,7 @@ def __getattr__(self, name):
     return getattr(self.wrapped, name)
 
   def apply(self, updates, params):
-    self.wrapped.apply_gradients(zip(updates, params))
+    self.wrapped.apply_gradients(zip(updates, params))  # pyrefly: ignore[missing-attribute]
 
 
 def is_tf_optimizer(optimizer):

sonnet/src/optimizers/optimizer_utils.py

@@ -52,7 +52,7 @@ def check_updates_parameters(updates: Sequence[types.ParameterUpdate],
 
 
 def check_same_dtype(update: types.ParameterUpdate, parameter: tf.Variable):
-  if update.dtype != parameter.dtype:
+  if update.dtype != parameter.dtype:  # pyrefly: ignore[missing-attribute]
     raise ValueError(
         "DType of update {!r} is not equal to that of parameter {!r}".format(
             update, parameter))