From 4a3bae0fe2837f68ef74025952ded3074e20592a Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Tue, 7 Jul 2026 15:30:28 -0700 Subject: [PATCH] Automated Code Change PiperOrigin-RevId: 944126193 --- sonnet/__init__.py | 2 +- sonnet/src/base.py | 14 +++++------ sonnet/src/batch_apply.py | 2 +- sonnet/src/bias.py | 10 ++++---- sonnet/src/conformance/descriptors.py | 2 +- sonnet/src/conformance/goldens.py | 2 +- sonnet/src/conv.py | 2 +- sonnet/src/conv_transpose.py | 10 ++++---- sonnet/src/custom_getter.py | 2 +- sonnet/src/embed.py | 2 +- sonnet/src/functional/haiku.py | 6 ++--- sonnet/src/functional/jax.py | 2 +- sonnet/src/initializers.py | 32 ++++++++++++------------ sonnet/src/metrics.py | 2 +- sonnet/src/nets/dnc/control.py | 4 +-- sonnet/src/nets/resnet.py | 8 +++--- sonnet/src/optimizers/optimizer_tests.py | 2 +- sonnet/src/optimizers/optimizer_utils.py | 2 +- sonnet/src/recurrent.py | 24 +++++++++--------- sonnet/src/reshape.py | 2 +- sonnet/src/utils.py | 10 ++++---- 21 files changed, 71 insertions(+), 71 deletions(-) diff --git a/sonnet/__init__.py b/sonnet/__init__.py index 629a1a2e..f2eecdbf 100644 --- a/sonnet/__init__.py +++ b/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 diff --git a/sonnet/src/base.py b/sonnet/src/base.py index d845014f..3eb7b454 100644 --- a/sonnet/src/base.py +++ b/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): diff --git a/sonnet/src/batch_apply.py b/sonnet/src/batch_apply.py index e6a6f0f4..f68f8954 100644 --- a/sonnet/src/batch_apply.py +++ b/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 diff --git a/sonnet/src/bias.py b/sonnet/src/bias.py index ff29c87d..5cfd2bde 100644 --- a/sonnet/src/bias.py +++ b/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): diff --git a/sonnet/src/conformance/descriptors.py b/sonnet/src/conformance/descriptors.py index 985ff711..463af8e1 100644 --- a/sonnet/src/conformance/descriptors.py +++ b/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: diff --git a/sonnet/src/conformance/goldens.py b/sonnet/src/conformance/goldens.py index 41ce4a28..be0741c7 100644 --- a/sonnet/src/conformance/goldens.py +++ b/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): diff --git a/sonnet/src/conv.py b/sonnet/src/conv.py index 1e6576cc..a6113c9b 100644 --- a/sonnet/src/conv.py +++ b/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 diff --git a/sonnet/src/conv_transpose.py b/sonnet/src/conv_transpose.py index d2ba2fb1..21cd7eba 100644 --- a/sonnet/src/conv_transpose.py +++ b/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: diff --git a/sonnet/src/custom_getter.py b/sonnet/src/custom_getter.py index e8e3def9..daf71ad4 100644 --- a/sonnet/src/custom_getter.py +++ b/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) diff --git a/sonnet/src/embed.py b/sonnet/src/embed.py index e87a5004..b9cb1e17 100644 --- a/sonnet/src/embed.py +++ b/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) diff --git a/sonnet/src/functional/haiku.py b/sonnet/src/functional/haiku.py index b992ea29..45c886c9 100644 --- a/sonnet/src/functional/haiku.py +++ b/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 diff --git a/sonnet/src/functional/jax.py b/sonnet/src/functional/jax.py index 5ac8a0e3..cd318e11 100644 --- a/sonnet/src/functional/jax.py +++ b/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 diff --git a/sonnet/src/initializers.py b/sonnet/src/initializers.py index a5144917..e4cf366f 100644 --- a/sonnet/src/initializers.py +++ b/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 diff --git a/sonnet/src/metrics.py b/sonnet/src/metrics.py index 9e438350..327dc43c 100644 --- a/sonnet/src/metrics.py +++ b/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 diff --git a/sonnet/src/nets/dnc/control.py b/sonnet/src/nets/dnc/control.py index 85f77059..2e8f0c25 100644 --- a/sonnet/src/nets/dnc/control.py +++ b/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) diff --git a/sonnet/src/nets/resnet.py b/sonnet/src/nets/resnet.py index d7254e42..08760068 100644 --- a/sonnet/src/nets/resnet.py +++ b/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") diff --git a/sonnet/src/optimizers/optimizer_tests.py b/sonnet/src/optimizers/optimizer_tests.py index ceb1e724..aea930e9 100644 --- a/sonnet/src/optimizers/optimizer_tests.py +++ b/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): diff --git a/sonnet/src/optimizers/optimizer_utils.py b/sonnet/src/optimizers/optimizer_utils.py index d904e2de..fc1df135 100644 --- a/sonnet/src/optimizers/optimizer_utils.py +++ b/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)) diff --git a/sonnet/src/recurrent.py b/sonnet/src/recurrent.py index 87a6e38a..96beb521 100644 --- a/sonnet/src/recurrent.py +++ b/sonnet/src/recurrent.py @@ -37,7 +37,7 @@ # pylint: enable=g-direct-tensorflow-import -class RNNCore(base.Module, metaclass=abc.ABCMeta): +class RNNCore(base.Module, metaclass=abc.ABCMeta): # pyrefly: ignore[invalid-inheritance] """Base class for Recurrent Neural Network cores. This class defines the basic functionality that every core should @@ -81,7 +81,7 @@ def initial_state(self, batch_size: types.IntegerLike, **kwargs): """ -class UnrolledRNN(base.Module, metaclass=abc.ABCMeta): +class UnrolledRNN(base.Module, metaclass=abc.ABCMeta): # pyrefly: ignore[invalid-inheritance] """Base class for unrolled Recurrent Neural Networks. This class is a generalization of :class:`RNNCore` which operates on @@ -528,7 +528,7 @@ def __call__(self, inputs: types.TensorNest, # For VanillaRNN, the next state of the RNN is the same as the outputs. return outputs, outputs - def initial_state(self, batch_size: int) -> tf.Tensor: + def initial_state(self, batch_size: int) -> tf.Tensor: # pyrefly: ignore[bad-override] """See base class.""" return tf.zeros(shape=[batch_size, self._hidden_size], dtype=self._dtype) @@ -832,7 +832,7 @@ def __call__(self, inputs, prev_state): return _lstm_fn(inputs, prev_state, self._w_i, self._w_h, self.b, self.projection) - def initial_state(self, batch_size: int) -> LSTMState: + def initial_state(self, batch_size: int) -> LSTMState: # pyrefly: ignore[bad-override] """See base class.""" return LSTMState( hidden=tf.zeros([batch_size, self._eff_hidden_size], dtype=self._dtype), @@ -946,7 +946,7 @@ def __call__(self, input_sequence, initial_state): return _specialized_unrolled_lstm(input_sequence, initial_state, self._w_i, self._w_h, self.b) - def initial_state(self, batch_size): + def initial_state(self, batch_size): # pyrefly: ignore[bad-override] """See base class.""" return LSTMState( hidden=tf.zeros([batch_size, self._hidden_size], dtype=self._dtype), @@ -1284,7 +1284,7 @@ def __init__(self, """ super().__init__(name) self._num_spatial_dims = num_spatial_dims - self._input_shape = list(input_shape) + self._input_shape = list(input_shape) # pyrefly: ignore[bad-argument-type] self._channel_index = 1 if (data_format is not None and data_format.startswith("NC")) else -1 self._output_channels = output_channels @@ -1336,7 +1336,7 @@ def input_to_hidden(self): def hidden_to_hidden(self): return self._hidden_to_hidden.w - def initial_state(self, batch_size): + def initial_state(self, batch_size): # pyrefly: ignore[bad-override] """See base class.""" shape = list(self._input_shape) shape[self._channel_index] = self._output_channels @@ -1355,7 +1355,7 @@ def _initialize(self, inputs): class Conv1DLSTM(_ConvNDLSTM): # pylint: disable=missing-docstring,empty-docstring - __doc__ = _ConvNDLSTM.__doc__.replace("``num_spatial_dims``", "1") + __doc__ = _ConvNDLSTM.__doc__.replace("``num_spatial_dims``", "1") # pyrefly: ignore[missing-attribute] def __init__(self, input_shape: types.ShapeLike, @@ -1406,7 +1406,7 @@ def __init__(self, class Conv2DLSTM(_ConvNDLSTM): # pylint: disable=missing-docstring,empty-docstring - __doc__ = _ConvNDLSTM.__doc__.replace("``num_spatial_dims``", "2") + __doc__ = _ConvNDLSTM.__doc__.replace("``num_spatial_dims``", "2") # pyrefly: ignore[missing-attribute] def __init__(self, input_shape: types.ShapeLike, @@ -1457,7 +1457,7 @@ def __init__(self, class Conv3DLSTM(_ConvNDLSTM): # pylint: disable=missing-docstring,empty-docstring - __doc__ = _ConvNDLSTM.__doc__.replace("``num_spatial_dims``", "3") + __doc__ = _ConvNDLSTM.__doc__.replace("``num_spatial_dims``", "3") # pyrefly: ignore[missing-attribute] def __init__(self, input_shape: types.ShapeLike, @@ -1584,7 +1584,7 @@ def __call__(self, inputs, prev_state): next_state = (1 - z) * prev_state + z * a return next_state, next_state - def initial_state(self, batch_size): + def initial_state(self, batch_size): # pyrefly: ignore[bad-override] """See base class.""" return tf.zeros([batch_size, self._hidden_size], dtype=self._dtype) @@ -1710,7 +1710,7 @@ def input_to_hidden(self): def hidden_to_hidden(self): return self._w_h - def initial_state(self, batch_size): + def initial_state(self, batch_size): # pyrefly: ignore[bad-override] """See base class.""" return tf.zeros([batch_size, self._hidden_size], dtype=self._dtype) diff --git a/sonnet/src/reshape.py b/sonnet/src/reshape.py index d03abc73..bef9d8e6 100644 --- a/sonnet/src/reshape.py +++ b/sonnet/src/reshape.py @@ -143,7 +143,7 @@ def __call__(self, inputs: tf.Tensor) -> tf.Tensor: self._initialize(inputs) # Resolve the wildcard if any. - output_shape = tuple(self._output_shape) + output_shape = tuple(self._output_shape) # pyrefly: ignore[bad-argument-type] if -1 in output_shape: reshaped_shape = inputs.shape[self._preserve_dims:] if reshaped_shape.is_fully_defined(): diff --git a/sonnet/src/utils.py b/sonnet/src/utils.py index 07a9b7f8..eba18397 100644 --- a/sonnet/src/utils.py +++ b/sonnet/src/utils.py @@ -36,11 +36,11 @@ def replicate( ) -> Tuple[T]: """Replicates entry in `element` `num_times` if needed.""" if not isinstance(element, collections.abc.Sequence): - return (element,) * num_times + return (element,) * num_times # pyrefly: ignore[bad-return] elif len(element) == 1: - return tuple(element * num_times) + return tuple(element * num_times) # pyrefly: ignore[bad-argument-type, unsupported-operation] elif len(element) == num_times: - return tuple(element) + return tuple(element) # pyrefly: ignore[bad-return] raise TypeError( "{} must be a scalar or sequence of length 1 or sequence of length {}." .format(name, num_times)) @@ -70,7 +70,7 @@ def _decorate_object(*args, **kwargs): @functools.wraps(f) def _decorate_bound_method(*args, **kwargs): - return decorator_fn(f, f.__self__, args, kwargs) + return decorator_fn(f, f.__self__, args, kwargs) # pyrefly: ignore[bad-argument-type] return _decorate_bound_method @@ -92,7 +92,7 @@ def _decorate_fn(*args, **kwargs): return _decorate_fn - return _decorator + return _decorator # pyrefly: ignore[bad-return] _SPATIAL_CHANNELS_FIRST = re.compile("^NC[^C]*$")