Automated Code Change

tensorflow_privacy/privacy/analysis/compute_noise_from_budget_lib.py

@@ -47,7 +47,7 @@ def make_accountant():
 
   target_noise = dp_accounting.calibrate_dp_mechanism(
       make_accountant, make_event_from_noise, target_epsilon, delta,
-      dp_accounting.LowerEndpointAndGuess(noise_lbd, noise_lbd * 2))
+      dp_accounting.LowerEndpointAndGuess(noise_lbd, noise_lbd * 2))  # pyrefly: ignore[bad-argument-count]
 
   print(
       'DP-SGD with sampling rate = {:.3g}% and noise_multiplier = {} iterated'

tensorflow_privacy/privacy/analysis/tree_aggregation_accountant.py

@@ -121,9 +121,9 @@ def compute_rdp_tree_restart(
         f"{steps_list}.")
 
   if np.isscalar(steps_list):
-    steps_list = [steps_list]
+    steps_list = [steps_list]  # pyrefly: ignore[bad-assignment]
 
-  for steps in steps_list:
+  for steps in steps_list:  # pyrefly: ignore[not-iterable]
     if steps < 0:
       raise ValueError(f"Steps must be non-negative, got {steps_list}")
 
@@ -132,7 +132,7 @@ def compute_rdp_tree_restart(
   else:
     rdp = np.array([
         _compute_rdp_tree_restart(noise_multiplier, steps_list, alpha)
-        for alpha in orders
+        for alpha in orders  # pyrefly: ignore[not-iterable]
     ])
 
   return rdp
@@ -317,11 +317,11 @@ def compute_rdp_single_tree(
       max_participation, min_separation, total_steps)
   if np.isscalar(orders):
     rdp = _compute_gaussian_rdp(noise_multiplier, sum_sensitivity_square,
-                                orders)
+                                orders)  # pyrefly: ignore[bad-argument-type]
   else:
     rdp = np.array([
         _compute_gaussian_rdp(noise_multiplier, sum_sensitivity_square, alpha)
-        for alpha in orders
+        for alpha in orders  # pyrefly: ignore[not-iterable]
     ])
   return rdp
 

tensorflow_privacy/privacy/dp_query/gaussian_query.py

@@ -88,7 +88,7 @@ def add_noise(v):
       random_normal = tf.random_normal_initializer(stddev=global_state.stddev)
 
       def add_noise(v):
-        return v + tf.cast(random_normal(tf.shape(input=v)), dtype=v.dtype)
+        return v + tf.cast(random_normal(tf.shape(input=v)), dtype=v.dtype)  # pyrefly: ignore[not-callable]
 
     result = tf.nest.map_structure(add_noise, sample_state)
     noise_multiplier = global_state.stddev / global_state.l2_norm_clip

tensorflow_privacy/privacy/dp_query/nested_query.py

@@ -101,7 +101,7 @@ def get_noised_result(self, sample_state, global_state):
 
     return (tf.nest.pack_sequence_as(self._queries, flat_estimates),
             tf.nest.pack_sequence_as(self._queries, flat_new_global_states),
-            dp_accounting.ComposedDpEvent(events=flat_events))
+            dp_accounting.ComposedDpEvent(events=flat_events))  # pyrefly: ignore[bad-argument-type]
 
   def derive_metrics(self, global_state):
     """Implements `tensorflow_privacy.DPQuery.derive_metrics`."""

tensorflow_privacy/privacy/dp_query/tree_aggregation.py

@@ -192,7 +192,7 @@ class TreeState(NamedTuple):
 def get_step_idx(state: TreeState) -> tf.Tensor:
   """Returns the current leaf node index based on `TreeState.level_buffer_idx`."""
   step_idx = tf.constant(-1, dtype=tf.int32)
-  for i in tf.range(len(state.level_buffer_idx)):
+  for i in tf.range(len(state.level_buffer_idx)):  # pyrefly: ignore[bad-argument-type]
     step_idx += tf.math.pow(2, state.level_buffer_idx[i])
   return step_idx
 
@@ -317,7 +317,7 @@ def get_cumsum_and_update(self,
     # noise. To update the buffer, let us find the lowest level that will switch
     # from a right child (not in the buffer) to a left child.
     level_idx = 0  # new leaf node starts from level 0
-    while tf.less(level_idx, len(level_buffer_idx)) and tf.equal(
+    while tf.less(level_idx, len(level_buffer_idx)) and tf.equal(  # pyrefly: ignore[bad-argument-type]
         level_idx, level_buffer_idx[level_idx]):
       level_idx += 1
     # Left child nodes for the level lower than `level_idx` will be removed
@@ -334,7 +334,7 @@ def get_cumsum_and_update(self,
     # i.e., `level_buffer_idx[level_idx] != level_idx`. Rename parameter to
     # buffer index for clarity.
     buffer_idx = level_idx
-    while tf.less(buffer_idx, len(level_buffer_idx)):
+    while tf.less(buffer_idx, len(level_buffer_idx)):  # pyrefly: ignore[bad-argument-type]
       new_level_buffer_idx = new_level_buffer_idx.write(
           write_buffer_idx, level_buffer_idx[buffer_idx])
       new_level_buffer = tf.nest.map_structure(
@@ -484,7 +484,7 @@ def get_cumsum_and_update(self,
     level_idx = 0  # new leaf node starts from level 0
     new_value, value_generator_state = self.value_generator.next(
         value_generator_state)
-    while tf.less(level_idx, len(level_buffer_idx)) and tf.equal(
+    while tf.less(level_idx, len(level_buffer_idx)) and tf.equal(  # pyrefly: ignore[bad-argument-type]
         level_idx, level_buffer_idx[level_idx]):
       # Recursively update if the current node is a right child.
       node_value, value_generator_state = self.value_generator.next(
@@ -504,7 +504,7 @@ def get_cumsum_and_update(self,
     # i.e., `level_buffer_idx[level_idx] != level_idx`. Rename parameter to
     # buffer index for clarity.
     buffer_idx = level_idx
-    while tf.less(buffer_idx, len(level_buffer_idx)):
+    while tf.less(buffer_idx, len(level_buffer_idx)):  # pyrefly: ignore[bad-argument-type]
       new_level_buffer_idx = new_level_buffer_idx.write(
           write_buffer_idx, level_buffer_idx[buffer_idx])
       new_level_buffer = tf.nest.map_structure(

tensorflow_privacy/privacy/dp_query/tree_range_query.py

@@ -284,6 +284,6 @@ def add_noise(v):
     random_normal = tf.random_normal_initializer(stddev=stddev, seed=seed)
 
     def add_noise(v):
-      return v + tf.cast(random_normal(tf.shape(input=v)), dtype=v.dtype)
+      return v + tf.cast(random_normal(tf.shape(input=v)), dtype=v.dtype)  # pyrefly: ignore[not-callable]
 
   return add_noise

tensorflow_privacy/privacy/estimators/test_utils.py

@@ -82,7 +82,7 @@ def make_input_fn(features, labels, training, batch_size=16):
   def input_fn():
     """An input function for training or evaluating."""
     # Convert the inputs to a Dataset.
-    dataset = tf.data.Dataset.from_tensor_slices((dict(features), labels))
+    dataset = tf.data.Dataset.from_tensor_slices((dict(features), labels))  # pyrefly: ignore[bad-argument-type]
 
     # Shuffle if in training mode.
     if training:
@@ -100,10 +100,10 @@ def model_fn(features, labels, mode, params, config=None):  # pylint: disable=un
     feature_layer = tf.keras.layers.DenseFeatures(feature_columns)
     inputs = feature_layer(features)
     hidden_layer = tf.keras.layers.Dense(units=3, activation='relu')
-    hidden_layer_values = hidden_layer(inputs)
+    hidden_layer_values = hidden_layer(inputs)  # pyrefly: ignore[not-callable]
     logits_layer = tf.keras.layers.Dense(
         units=head.logits_dimension, activation=None)
-    logits = logits_layer(hidden_layer_values)
+    logits = logits_layer(hidden_layer_values)  # pyrefly: ignore[not-callable]
     return head.create_estimator_spec(
         features=features,
         labels=labels,

tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads.py

@@ -69,7 +69,7 @@ def _compute_gradient_norms_internal(
   if trainable_vars is not None:
     # Create a set using `ref()` for fast set membership check. tf.Variable
     # itself is not hashable.
-    trainable_vars = set([v.ref() for v in trainable_vars])
+    trainable_vars = set([v.ref() for v in trainable_vars])  # pyrefly: ignore[bad-assignment]
 
   layer_sqr_norm_fns = collections.defaultdict(list)
   # The case of shared weights:
@@ -310,7 +310,7 @@ def compute_clipped_gradients_and_outputs(
     # (e.g. mean_squared_error and binary_crossentropy). However it
     # is not defined in the contract so may not hold, especially for
     # custom losses.
-    y_pred = input_model(x_batch, training=True)
+    y_pred = input_model(x_batch, training=True)  # pyrefly: ignore[not-callable]
     mb_y_batch = common_manip_utils.maybe_add_microbatch_axis(
         y_batch, num_microbatches
     )

tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py

@@ -61,9 +61,9 @@ def test_loss_fn(
 ) -> tf.Tensor:
   # Define a loss function which is unlikely to be coincidently defined.
   if weights is None:
-    weights = 1.0
+    weights = 1.0  # pyrefly: ignore[bad-assignment]
   loss = 3.14 * tf.reduce_sum(
-      tf.cast(weights, tf.float32) * tf.square(x - y), axis=1
+      tf.cast(weights, tf.float32) * tf.square(x - y), axis=1  # pyrefly: ignore[unsupported-operation]
   )
   return loss
 
@@ -83,7 +83,7 @@ def compute_true_gradient_norms(
     loss_config['reduction'] = tf.keras.losses.Reduction.NONE
     per_example_loss_fn = input_model.loss.from_config(loss_config)
   with tf.GradientTape(persistent=True) as tape:
-    y_pred = input_model(x_batch)
+    y_pred = input_model(x_batch)  # pyrefly: ignore[not-callable]
     loss = per_example_loss_fn(y_batch, y_pred, weight_batch)
     if num_microbatches is not None:
       loss = tf.reduce_mean(
@@ -143,7 +143,7 @@ def get_computed_and_true_norms_from_model(
       trainable_vars = l.trainable_variables
       if trainable_vars:
         break
-  y_pred = model(x_batch)
+  y_pred = model(x_batch)  # pyrefly: ignore[not-callable]
   y_batch = tf.ones_like(y_pred)
   tf.keras.utils.set_random_seed(rng_seed)
   computed_norms = clip_grads.compute_gradient_norms(
@@ -266,7 +266,7 @@ def make_two_layer_functional_model(
   inputs = tf.keras.Input(shape=input_dims)
   layer1 = layer_generator(input_dims, output_dims)
   temp1 = layer1(inputs)
-  temp2 = tf.keras.layers.Dense(1)(temp1)
+  temp2 = tf.keras.layers.Dense(1)(temp1)  # pyrefly: ignore[not-callable]
   outputs = reshape_and_sum(temp2)
   return tf.keras.Model(inputs=inputs, outputs=outputs)
 
@@ -284,7 +284,7 @@ def make_two_tower_model(
   layer2 = layer_generator(input_dims, output_dims)
   temp2 = layer2(inputs2)
   temp3 = tf.add(temp1, temp2)
-  temp4 = tf.keras.layers.Dense(1)(temp3)
+  temp4 = tf.keras.layers.Dense(1)(temp3)  # pyrefly: ignore[not-callable]
   outputs = reshape_and_sum(temp4)
   return tf.keras.Model(inputs=[inputs1, inputs2], outputs=outputs)
 
@@ -336,7 +336,7 @@ def make_dense_bow_model(
   example_embs = tf.expand_dims(
       tf.reduce_sum(feature_embs, axis=reduction_axes), axis=-1
   )
-  outputs = tf.keras.layers.Dense(1)(example_embs)
+  outputs = tf.keras.layers.Dense(1)(example_embs)  # pyrefly: ignore[not-callable]
   return tf.keras.Model(inputs=inputs, outputs=outputs)
 
 
@@ -356,7 +356,7 @@ def make_weighted_bow_model(
     raise ValueError('Expected `output_dims` to be of size 1.')
   feature_embs = emb_layer(inputs)
   # Use deterministic weights to avoid seeding issues on TPUs.
-  feature_shape = input_dims + output_dims
+  feature_shape = input_dims + output_dims  # pyrefly: ignore[unsupported-operation]
   feature_weights = tf.expand_dims(
       tf.reshape(
           tf.range(np.prod(feature_shape), dtype=tf.float32),
@@ -373,5 +373,5 @@ def make_weighted_bow_model(
   example_embs = tf.expand_dims(
       tf.reduce_sum(weighted_embs, axis=reduction_axes), axis=-1
   )
-  outputs = tf.keras.layers.Dense(1)(example_embs)
+  outputs = tf.keras.layers.Dense(1)(example_embs)  # pyrefly: ignore[not-callable]
   return tf.keras.Model(inputs=inputs, outputs=outputs)

tensorflow_privacy/privacy/fast_gradient_clipping/gradient_clipping_utils.py

@@ -242,7 +242,7 @@ def model_forward_backward_pass(
   if trainable_vars is not None:
     # Create a set using `ref()` for fast set membership check. tf.Variable
     # itself is not hashable.
-    trainable_vars = set([v.ref() for v in trainable_vars])
+    trainable_vars = set([v.ref() for v in trainable_vars])  # pyrefly: ignore[bad-assignment]
   layer_vars = collections.defaultdict(list)
   for registry_fn_output in filtered_outputs:
     if trainable_vars is None or any(

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense.py

@@ -68,7 +68,7 @@ def dense_layer_computation(
     raise ValueError("Only layer inputs of length 1 are permitted.")
   orig_activation = layer_instance.activation
   layer_instance.activation = None
-  base_vars = layer_instance(*input_args)
+  base_vars = layer_instance(*input_args)  # pyrefly: ignore[not-callable]
   tape.watch(base_vars)
   layer_instance.activation = orig_activation
   outputs = orig_activation(base_vars) if orig_activation else base_vars

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense.py

@@ -53,7 +53,7 @@ def einsum_layer_computation(
   # To avoid this case, we watch the variables that are only generated by the
   # linear transformation of the `EinsumDense` layer instance.
   layer_instance.activation = None
-  base_vars = layer_instance(*input_args)
+  base_vars = layer_instance(*input_args)  # pyrefly: ignore[not-callable]
   tape.watch(base_vars)
   layer_instance.activation = orig_activation
   outputs = orig_activation(base_vars) if orig_activation else base_vars

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding.py

@@ -63,7 +63,7 @@ def embedding_layer_computation(
           "The experimental embedding feature"
           "'_use_one_hot_matmul' is not supported."
       )
-  input_ids = tf.cast(*input_args, tf.int32)
+  input_ids = tf.cast(*input_args, tf.int32)  # pyrefly: ignore[bad-argument-count]
   if len(layer_instance.trainable_variables) != 1:
     raise ValueError(
         "Only layer instances with only one set of trainable variables"

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/multi_head_attention.py

@@ -116,9 +116,9 @@ def multi_head_attention_layer_computation(
     key = key.to_tensor(shape=bounding_shape)
     value = value.to_tensor(shape=bounding_shape)
   elif key_is_ragged:
-    key = key.to_tensor(shape=tf.shape(value))
+    key = key.to_tensor(shape=tf.shape(value))  # pyrefly: ignore[missing-attribute]
   elif value_is_ragged:
-    value = value.to_tensor(shape=tf.shape(key))
+    value = value.to_tensor(shape=tf.shape(key))  # pyrefly: ignore[missing-attribute]
   else:
     pass
   # ------------------------------

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding.py

@@ -51,15 +51,15 @@ def nlp_on_device_embedding_layer_computation(
       )
   # NOTE: Since the implementation of `tfm.nlp.layers.OnDeviceEmbedding` uses
   # `.set_shape()`, we can assume that inputs are not ragged.
-  input_ids = tf.cast(*input_args, tf.int32)
+  input_ids = tf.cast(*input_args, tf.int32)  # pyrefly: ignore[bad-argument-count]
   if len(layer_instance.trainable_variables) != 1:
     raise ValueError(
         "Only layer instances with only one set of trainable variables"
         "are permitted."
     )
   base_vars = layer_instance.trainable_variables[0]
   tape.watch(base_vars)
-  outputs = layer_instance(input_ids)
+  outputs = layer_instance(input_ids)  # pyrefly: ignore[not-callable]
 
   def sqr_norm_fn(base_vars_grads: tf.IndexedSlices):
     return registry_function_utils.embedding_sqr_norm_fn(

tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_position_embedding.py

@@ -44,8 +44,8 @@ def nlp_position_embedding_layer_computation(
   del input_kwargs
   if len(input_args) != 1:
     raise ValueError("Only layer inputs of length 1 are permitted.")
-  input_ids = tf.cast(*input_args, tf.int32)
-  base_vars = layer_instance(input_ids)
+  input_ids = tf.cast(*input_args, tf.int32)  # pyrefly: ignore[bad-argument-count]
+  base_vars = layer_instance(input_ids)  # pyrefly: ignore[not-callable]
   tape.watch(base_vars)
 
   def sqr_norm_fn(grads):

tensorflow_privacy/privacy/keras_models/dp_keras_model.py

@@ -223,9 +223,9 @@ def _compute_per_example_grads(self, microbatched_data):
           tf.keras.utils.unpack_x_y_sample_weight(microbatched_data)
       )
       with tf.GradientTape() as tape:
-        microbatched_y_pred = self(microbatched_x, training=True)
+        microbatched_y_pred = self(microbatched_x, training=True)  # pyrefly: ignore[not-callable]
         # NOTE: `self._clipping_loss` does not include any regularization terms.
-        microbatched_loss = self._clipping_loss(
+        microbatched_loss = self._clipping_loss(  # pyrefly: ignore[not-callable]
             microbatched_y,
             microbatched_y_pred,
             sample_weight=microbatched_weights,
@@ -410,7 +410,7 @@ def reshape_fn(z):
             self._compute_per_example_grads,
             microbatched_data,
         )
-        y_pred = self(x, training=True)
+        y_pred = self(x, training=True)  # pyrefly: ignore[not-callable]
         grads = tf.nest.map_structure(
             self._reduce_per_example_grads, clipped_grads
         )

tensorflow_privacy/privacy/logistic_regression/multinomial_logistic.py

@@ -183,7 +183,7 @@ def make_event_from_param(noise_multiplier):
       make_event_from_param,
       epsilon,
       delta,
-      dp_accounting.LowerEndpointAndGuess(0, 1),
+      dp_accounting.LowerEndpointAndGuess(0, 1),  # pyrefly: ignore[bad-argument-count]
       tol=tolerance)
 
 

tensorflow_privacy/privacy/optimizers/clip_and_aggregate_gradients.py

@@ -73,7 +73,7 @@ def _expand_dims(e, v):
         axis=0)
     return tf.reshape(e, new_shape)
 
-  return [
+  return [  # pyrefly: ignore[bad-return]
       v *
       _expand_dims(tf.cast(clip_ratio, v.dtype), v) if v is not None else None
       for v in vals
@@ -219,7 +219,7 @@ def loop_fn(i):
 
   if normalize or l2_norm_clip is not None:
     values, indices, dense_shape = zip(*grads)
-    values = _batch_clip_by_global_norm(values, normalize, l2_norm_clip)
+    values = _batch_clip_by_global_norm(values, normalize, l2_norm_clip)  # pyrefly: ignore[bad-argument-type]
     grads = zip(values, indices, dense_shape)
 
   new_output = []

tensorflow_privacy/privacy/optimizers/dp_optimizer.py

@@ -34,7 +34,7 @@ def make_optimizer_class(cls):
   if has_compute_gradients:
     child_code = cls.compute_gradients.__code__
   GATE_OP = tf.compat.v1.train.Optimizer.GATE_OP  # pylint: disable=invalid-name
-  if has_compute_gradients and child_code is not parent_code:
+  if has_compute_gradients and child_code is not parent_code:  # pyrefly: ignore[unbound-name]
     logging.warning(
         'WARNING: Calling make_optimizer_class() on class %s that overrides '
         'method compute_gradients(). Check to ensure that '
@@ -278,7 +278,7 @@ def make_gaussian_optimizer_class(cls):
     A subclass of `cls` using DP-SGD with Gaussian averaging.
   """
 
-  class DPGaussianOptimizerClass(make_optimizer_class(cls)):  # pylint: disable=missing-class-docstring
+  class DPGaussianOptimizerClass(make_optimizer_class(cls)):  # pylint: disable=missing-class-docstring  # pyrefly: ignore[invalid-inheritance]
     __doc__ = ("""DP subclass of `{}`.
 
        You can use this as a differentially private replacement for

tensorflow_privacy/privacy/optimizers/dp_optimizer_keras_sparse.py

@@ -222,7 +222,7 @@ def _prepare_local(self, var_device, var_dtype, apply_state):
       if self.gradient_accumulation_steps > 1:
         apply_update = tf.math.equal(
             tf.math.floormod(
-                self._acc_iterations + 1, self.gradient_accumulation_steps
+                self._acc_iterations + 1, self.gradient_accumulation_steps  # pyrefly: ignore[unsupported-operation]
             ),
             0,
         )