From e5d2dea72322d055887e7961dc44205da0540c8c Mon Sep 17 00:00:00 2001
From: Hana Joo <hanajoo@google.com>
Date: Wed, 8 Jul 2026 01:55:39 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 944355784
---
 .../analysis/compute_noise_from_budget_lib.py |  2 +-
 .../analysis/tree_aggregation_accountant.py   | 10 ++++-----
 .../privacy/dp_query/gaussian_query.py        |  2 +-
 .../privacy/dp_query/nested_query.py          |  2 +-
 .../privacy/dp_query/tree_aggregation.py      | 10 ++++-----
 .../privacy/dp_query/tree_range_query.py      |  2 +-
 .../privacy/estimators/test_utils.py          |  6 +++---
 .../fast_gradient_clipping/clip_grads.py      |  4 ++--
 .../common_test_utils.py                      | 18 ++++++++--------
 .../gradient_clipping_utils.py                |  2 +-
 .../registry_functions/dense.py               |  2 +-
 .../registry_functions/einsum_dense.py        |  2 +-
 .../registry_functions/embedding.py           |  2 +-
 .../multi_head_attention.py                   |  4 ++--
 .../nlp_on_device_embedding.py                |  4 ++--
 .../nlp_position_embedding.py                 |  4 ++--
 .../privacy/keras_models/dp_keras_model.py    |  6 +++---
 .../multinomial_logistic.py                   |  2 +-
 .../clip_and_aggregate_gradients.py           |  4 ++--
 .../privacy/optimizers/dp_optimizer.py        |  4 ++--
 .../optimizers/dp_optimizer_keras_sparse.py   |  2 +-
 .../privacy_tests/epsilon_lower_bound.py      | 21 ++++++++++---------
 .../advanced_mia.py                           | 14 ++++++-------
 .../advanced_mia_example.py                   |  2 +-
 .../codelabs/example.py                       |  2 +-
 .../data_structures.py                        | 17 ++++++++-------
 .../membership_inference_attack.py            |  2 +-
 .../membership_inference_attack/plotting.py   |  4 ++--
 .../seq2seq_mia.py                            |  4 ++--
 .../tf_estimator_evaluation_example.py        | 10 ++++-----
 .../privacy/privacy_tests/utils.py            | 14 ++++++-------
 .../registry_functions/embedding.py           |  2 +-
 .../sparse_noise_utils.py                     |  6 +++---
 tutorials/lm_dpsgd_tutorial.py                |  8 +++----
 tutorials/mnist_dpsgd_tutorial_common.py      |  8 +++----
 tutorials/mnist_dpsgd_tutorial_eager.py       |  6 +++---
 tutorials/mnist_dpsgd_tutorial_vectorized.py  |  6 +++---
 tutorials/mnist_lr_tutorial.py                |  6 +++---
 tutorials/movielens_tutorial.py               | 10 ++++-----
 39 files changed, 119 insertions(+), 117 deletions(-)

diff --git a/tensorflow_privacy/privacy/analysis/compute_noise_from_budget_lib.py b/tensorflow_privacy/privacy/analysis/compute_noise_from_budget_lib.py
index e561130d..190266d6 100644
--- a/tensorflow_privacy/privacy/analysis/compute_noise_from_budget_lib.py
+++ b/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'
diff --git a/tensorflow_privacy/privacy/analysis/tree_aggregation_accountant.py b/tensorflow_privacy/privacy/analysis/tree_aggregation_accountant.py
index a7d95059..5cfa4845 100644
--- a/tensorflow_privacy/privacy/analysis/tree_aggregation_accountant.py
+++ b/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
 
diff --git a/tensorflow_privacy/privacy/dp_query/gaussian_query.py b/tensorflow_privacy/privacy/dp_query/gaussian_query.py
index 834abccc..5d03aedc 100644
--- a/tensorflow_privacy/privacy/dp_query/gaussian_query.py
+++ b/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
diff --git a/tensorflow_privacy/privacy/dp_query/nested_query.py b/tensorflow_privacy/privacy/dp_query/nested_query.py
index 235cf99f..f02eba91 100644
--- a/tensorflow_privacy/privacy/dp_query/nested_query.py
+++ b/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`."""
diff --git a/tensorflow_privacy/privacy/dp_query/tree_aggregation.py b/tensorflow_privacy/privacy/dp_query/tree_aggregation.py
index 436dcff3..586a395c 100644
--- a/tensorflow_privacy/privacy/dp_query/tree_aggregation.py
+++ b/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(
diff --git a/tensorflow_privacy/privacy/dp_query/tree_range_query.py b/tensorflow_privacy/privacy/dp_query/tree_range_query.py
index 8213fa41..8bd7ff00 100644
--- a/tensorflow_privacy/privacy/dp_query/tree_range_query.py
+++ b/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
diff --git a/tensorflow_privacy/privacy/estimators/test_utils.py b/tensorflow_privacy/privacy/estimators/test_utils.py
index a4bfe13b..1dbc10ed 100644
--- a/tensorflow_privacy/privacy/estimators/test_utils.py
+++ b/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,
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads.py b/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads.py
index be7d92f9..7ea2a6f4 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/clip_grads.py
+++ b/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
     )
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py b/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py
index 5977e664..0aa95fe2 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/common_test_utils.py
+++ b/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)
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/gradient_clipping_utils.py b/tensorflow_privacy/privacy/fast_gradient_clipping/gradient_clipping_utils.py
index 989a69c2..52f2f779 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/gradient_clipping_utils.py
+++ b/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(
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense.py
index 4218a1ad..20e1ffd6 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/dense.py
+++ b/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
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense.py
index 8c5ecbc7..74c67a49 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense.py
+++ b/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
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding.py
index 8d146461..c72f86b0 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/embedding.py
+++ b/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"
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/multi_head_attention.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/multi_head_attention.py
index 4078c064..25d02122 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/multi_head_attention.py
+++ b/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
   # ------------------------------
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding.py
index c4dc21db..724de537 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding.py
+++ b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_on_device_embedding.py
@@ -51,7 +51,7 @@ 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"
@@ -59,7 +59,7 @@ def nlp_on_device_embedding_layer_computation(
     )
   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(
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_position_embedding.py b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_position_embedding.py
index cf6a1d7c..939a4124 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/nlp_position_embedding.py
+++ b/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):
diff --git a/tensorflow_privacy/privacy/keras_models/dp_keras_model.py b/tensorflow_privacy/privacy/keras_models/dp_keras_model.py
index 481f4a93..e0567557 100644
--- a/tensorflow_privacy/privacy/keras_models/dp_keras_model.py
+++ b/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
         )
diff --git a/tensorflow_privacy/privacy/logistic_regression/multinomial_logistic.py b/tensorflow_privacy/privacy/logistic_regression/multinomial_logistic.py
index 522c9a27..1634a92c 100644
--- a/tensorflow_privacy/privacy/logistic_regression/multinomial_logistic.py
+++ b/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)
 
 
diff --git a/tensorflow_privacy/privacy/optimizers/clip_and_aggregate_gradients.py b/tensorflow_privacy/privacy/optimizers/clip_and_aggregate_gradients.py
index cd86c02a..4adf2e58 100644
--- a/tensorflow_privacy/privacy/optimizers/clip_and_aggregate_gradients.py
+++ b/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 = []
diff --git a/tensorflow_privacy/privacy/optimizers/dp_optimizer.py b/tensorflow_privacy/privacy/optimizers/dp_optimizer.py
index 10cb7c56..39b18210 100644
--- a/tensorflow_privacy/privacy/optimizers/dp_optimizer.py
+++ b/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
diff --git a/tensorflow_privacy/privacy/optimizers/dp_optimizer_keras_sparse.py b/tensorflow_privacy/privacy/optimizers/dp_optimizer_keras_sparse.py
index 1cc49803..96887c69 100644
--- a/tensorflow_privacy/privacy/optimizers/dp_optimizer_keras_sparse.py
+++ b/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,
         )
diff --git a/tensorflow_privacy/privacy/privacy_tests/epsilon_lower_bound.py b/tensorflow_privacy/privacy/privacy_tests/epsilon_lower_bound.py
index b047e4a6..1684ec28 100644
--- a/tensorflow_privacy/privacy/privacy_tests/epsilon_lower_bound.py
+++ b/tensorflow_privacy/privacy/privacy_tests/epsilon_lower_bound.py
@@ -153,7 +153,7 @@ def __init__(self,
   def compute_epsilon_lower_bound(self,
                                   method: BoundMethod,
                                   k: Optional[int] = None
-                                 ) -> npt.NDArray[float]:
+                                 ) -> npt.NDArray[float]:  # pyrefly: ignore[bad-specialization]
     """Computes lower bound w/ a specified method and returns top-k epsilons.
 
     Args:
@@ -175,7 +175,7 @@ def compute_epsilon_lower_bound(self,
   def compute_epsilon_lower_bounds(
       self,
       methods: Optional[Iterable[BoundMethod]] = None,
-      k: Optional[int] = None) -> Dict[BoundMethod, npt.NDArray[float]]:
+      k: Optional[int] = None) -> Dict[BoundMethod, npt.NDArray[float]]:  # pyrefly: ignore[bad-specialization]
     """Computes lower bounds with all methods and returns the top-k epsilons.
 
     Args:
@@ -237,7 +237,8 @@ def __init__(self, tp: Union[Sequence[int], int], fp: Union[Sequence[int],
       self._is_scalar = True
     if isinstance(fp, numbers.Number):
       fp = [fp]
-    if len(tp) != len(fp):
+    if len(tp) != len(fp):  # pyrefly: ignore[bad-argument-type]
+      # pyrefly: ignore[bad-argument-type]
       raise ValueError('tp and fp should have the same number of elements, '
                        f'but get {len(tp)} and {len(fp)} respectively.')
     # Some methods need the original values.
@@ -276,7 +277,7 @@ def _get_statistics(self, tp, fp):
     return tpr, fpr, fnr, tnr
 
   def compute_bound(self,
-                    method: BoundMethod) -> Union[float, npt.NDArray[float]]:
+                    method: BoundMethod) -> Union[float, npt.NDArray[float]]:  # pyrefly: ignore[bad-specialization]
     """Computes ratio bound using a specified method.
 
     Args:
@@ -295,7 +296,7 @@ def compute_bound(self,
   def compute_bounds(
       self,
       methods: Optional[Iterable[BoundMethod]] = None
-  ) -> Dict[BoundMethod, Union[float, npt.NDArray[float]]]:
+  ) -> Dict[BoundMethod, Union[float, npt.NDArray[float]]]:  # pyrefly: ignore[bad-specialization]
     """Computes ratio bounds for specified methods.
 
     Args:
@@ -310,7 +311,7 @@ def compute_bounds(
         for method in methods or self.available_methods.keys()
     }
 
-  def _bound_katz_log(self) -> npt.NDArray[float]:
+  def _bound_katz_log(self) -> npt.NDArray[float]:  # pyrefly: ignore[bad-specialization]
     """Uses the logarithm Katz method to compute lower bound of ratio."""
     tp, fp = self._tp, np.where(self._idx_fp_0, 0.5, self._fp)
     tpr, fpr, fnr, tnr = self._get_statistics(tp, fp)
@@ -319,7 +320,7 @@ def _bound_katz_log(self) -> npt.NDArray[float]:
     return np.where(self._idx_tp_0, 0,
                     empirical_ratio * np.exp(self._z * sqrt_term))
 
-  def _bound_adjusted_log(self) -> npt.NDArray[float]:
+  def _bound_adjusted_log(self) -> npt.NDArray[float]:  # pyrefly: ignore[bad-specialization]
     """Uses the logarithm Walters method to compute lower bound of ratio."""
     log_empirical_ratio = (
         np.log((self._tp + 0.5) / (self._pos_size + 0.5)) - np.log(
@@ -330,7 +331,7 @@ def _bound_adjusted_log(self) -> npt.NDArray[float]:
         np.logical_and(self._idx_tp_0, self._idx_fp_0), 0,
         np.exp(log_empirical_ratio) * np.exp(self._z * sqrt_term))
 
-  def _bound_bailey(self) -> npt.NDArray[float]:
+  def _bound_bailey(self) -> npt.NDArray[float]:  # pyrefly: ignore[bad-specialization]
     """Uses the Bailey method to compute lower bound of ratio."""
     tp = np.where(self._tp_orig == self._pos_size, self._pos_size - 0.5,
                   self._tp_orig)
@@ -347,7 +348,7 @@ def _bound_bailey(self) -> npt.NDArray[float]:
         self._idx_tp_0, 0,
         empirical_ratio * (power_3_term_numer / power_3_term_denom)**3)
 
-  def _bound_inv_hyperbolic_sine(self) -> npt.NDArray[float]:
+  def _bound_inv_hyperbolic_sine(self) -> npt.NDArray[float]:  # pyrefly: ignore[bad-specialization]
     """Uses the inverse sinh method to compute lower bound of ratio."""
     tp, fp = self._tp, np.where(self._idx_fp_0, self._z**2, self._fp)
     empirical_ratio = (tp / fp) / (self._pos_size / self._neg_size)
@@ -356,7 +357,7 @@ def _bound_inv_hyperbolic_sine(self) -> npt.NDArray[float]:
     return np.where(self._idx_tp_0, 0,
                     empirical_ratio * np.exp(2 * np.arcsinh(in_inve_sinh)))
 
-  def _bound_clopper_pearson(self) -> npt.NDArray[float]:
+  def _bound_clopper_pearson(self) -> npt.NDArray[float]:  # pyrefly: ignore[bad-specialization]
     """Uses the Clopper-Pearson method to compute lower bound of ratio."""
     # proportion_confint uses alpha / 2 budget on upper and lower, so total
     # budget will be 2 * alpha/2 = alpha.
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia.py
index 332972a8..2c7c6114 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia.py
@@ -149,10 +149,10 @@ def compute_score_lira(stat_target: Union[np.ndarray, Sequence[float]],
     # standard deviation of statistics across shadow models and examples
     if option in ['in', 'both']:
       std_in = np.nanstd(
-          np.concatenate([l - m[np.newaxis] for l, m in zip(stat_in, avg_in)]))
+          np.concatenate([l - m[np.newaxis] for l, m in zip(stat_in, avg_in)]))  # pyrefly: ignore[unbound-name]
     if option in ['out', 'both']:
       std_out = np.nanstd(
-          np.concatenate([l - m[np.newaxis] for l, m in zip(stat_out, avg_out)
+          np.concatenate([l - m[np.newaxis] for l, m in zip(stat_out, avg_out)  # pyrefly: ignore[unbound-name]
                          ]))
   else:
     # standard deviation of statistics across shadow models
@@ -166,16 +166,16 @@ def compute_score_lira(stat_target: Union[np.ndarray, Sequence[float]],
 
   stat_target = np.array(stat_target)
   if option in ['in', 'both']:
-    log_pr_in = scipy.stats.norm.logpdf(stat_target, avg_in, std_in + 1e-30)
+    log_pr_in = scipy.stats.norm.logpdf(stat_target, avg_in, std_in + 1e-30)  # pyrefly: ignore[unbound-name]
   if option in ['out', 'both']:
-    log_pr_out = scipy.stats.norm.logpdf(stat_target, avg_out, std_out + 1e-30)
+    log_pr_out = scipy.stats.norm.logpdf(stat_target, avg_out, std_out + 1e-30)  # pyrefly: ignore[unbound-name]
 
   if option == 'both':
-    scores = -(log_pr_in - log_pr_out).mean(axis=1)
+    scores = -(log_pr_in - log_pr_out).mean(axis=1)  # pyrefly: ignore[unbound-name]
   elif option == 'in':
-    scores = -log_pr_in.mean(axis=1)
+    scores = -log_pr_in.mean(axis=1)  # pyrefly: ignore[unbound-name]
   else:
-    scores = log_pr_out.mean(axis=1)
+    scores = log_pr_out.mean(axis=1)  # pyrefly: ignore[unbound-name]
   return scores
 
 
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia_example.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia_example.py
index f033aa82..fc827e5c 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia_example.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/advanced_mia_example.py
@@ -137,7 +137,7 @@ def main(unused_argv):
     in_indices.append(np.random.binomial(1, 0.5, n).astype(bool))
 
     model = small_cnn()
-    if _MODEL_DIR.value and os.path.exists(model_path):  # Load if exists
+    if _MODEL_DIR.value and os.path.exists(model_path):  # Load if exists  # pyrefly: ignore[unbound-name]
       model(x[:1])  # use this to make the `load_weights` work
       model.load_weights(model_path)
       print(f'Loaded model #{i} with {in_indices[-1].sum()} examples.')
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/codelabs/example.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/codelabs/example.py
index 7c399c7d..59977b4c 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/codelabs/example.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/codelabs/example.py
@@ -203,7 +203,7 @@ def main(unused_argv):
   # Example of saving the results to the file and loading them back.
   with tempfile.TemporaryDirectory() as tmpdirname:
     filepath = os.path.join(tmpdirname, "results.pickle")
-    attack_results.save(filepath)
+    attack_results.save(filepath)  # pyrefly: ignore[unbound-name]
     loaded_results = data_structures.AttackResults.load(filepath)
     print(loaded_results.summary(by_slices=False))
 
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/data_structures.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/data_structures.py
index 76bbd941..11131837 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/data_structures.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/data_structures.py
@@ -73,7 +73,7 @@ def __str__(self):
       return 'Loss percentiles: %d-%d' % self.value
 
     if self.feature == SlicingFeature.CUSTOM:
-      custom_train_indices, custom_test_indices, slice_value, slice_name = (
+      custom_train_indices, custom_test_indices, slice_value, slice_name = (  # pyrefly: ignore[not-iterable]
           self.value
       )
       if slice_name is not None:
@@ -142,7 +142,8 @@ def __post_init__(self):
         self.all_custom_test_indices):
       raise ValueError('custom_train_indices and custom_test_indices must '
                        'be provided or set to None at the same time.')
-    if len(self.all_custom_train_indices) != len(self.all_custom_test_indices):
+    if len(self.all_custom_train_indices) != len(self.all_custom_test_indices):  # pyrefly: ignore[bad-argument-type]
+      # pyrefly: ignore[bad-argument-type]
       raise ValueError('all_custom_train_indices and all_custom_test_indices '
                        'should have the same length, but got'
                        f'{len(self.all_custom_train_indices)} and '
@@ -425,7 +426,7 @@ def get_entropy_train(self):
                                 'applicable for multi-label data.')
     if self.entropy_train is not None:
       return self.entropy_train
-    return self._get_entropy(self.logits_train, self.labels_train)
+    return self._get_entropy(self.logits_train, self.labels_train)  # pyrefly: ignore[bad-argument-type]
 
   def get_entropy_test(self):
     """Calculates prediction entropy for the test set."""
@@ -435,7 +436,7 @@ def get_entropy_test(self):
                                 'applicable for multi-label data.')
     if self.entropy_test is not None:
       return self.entropy_test
-    return self._get_entropy(self.logits_test, self.labels_test)
+    return self._get_entropy(self.logits_test, self.labels_test)  # pyrefly: ignore[bad-argument-type]
 
   def get_train_shape(self):
     """Returns the shape of the training set."""
@@ -1248,7 +1249,7 @@ def get_result_with_max_auc(self) -> Optional[SingleAttackResult]:
       logging.info(('Suspiciously low AUC detected: %.2f. '
                     'There might be a bug in the classifier'), min(aucs))
 
-    return self.single_attack_results[np.argmax(aucs)]
+    return self.single_attack_results[np.argmax(aucs)]  # pyrefly: ignore[bad-index]
 
   def get_result_with_max_attacker_advantage(
       self,
@@ -1256,7 +1257,7 @@ def get_result_with_max_attacker_advantage(
     """Get the result with maximum advantage for all attacks and slices."""
     if not self.single_attack_results:
       return None
-    return self.single_attack_results[np.argmax([
+    return self.single_attack_results[np.argmax([  # pyrefly: ignore[bad-index]
         result.get_attacker_advantage() for result in self.single_attack_results
     ])]
 
@@ -1264,7 +1265,7 @@ def get_result_with_max_ppv(self) -> Optional[SingleAttackResult]:
     """Gets the result with max positive predictive value for all attacks and slices."""
     if not self.single_attack_results:
       return None
-    return self.single_attack_results[np.argmax(
+    return self.single_attack_results[np.argmax(  # pyrefly: ignore[bad-index]
         [result.get_ppv() for result in self.single_attack_results])]
 
   def get_result_with_max_epsilon(self) -> SingleAttackResult:
@@ -1273,7 +1274,7 @@ def get_result_with_max_epsilon(self) -> SingleAttackResult:
         result.get_epsilon_lower_bound().mean()
         for result in self.single_attack_results
     ]
-    return self.single_attack_results[np.argmax(avg_epsilon_bounds)]
+    return self.single_attack_results[np.argmax(avg_epsilon_bounds)]  # pyrefly: ignore[bad-index]
 
   def save(self, filepath):
     """Saves self to a pickle file."""
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack.py
index 988a50da..d012f05b 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack.py
@@ -119,7 +119,7 @@ def _run_trained_attack(
   # Predict the left out with the last attacker
   if left_out_indices.size:
     assert np.all(np.isnan(scores[left_out_indices]))
-    scores[left_out_indices] = attacker.predict(features[left_out_indices])
+    scores[left_out_indices] = attacker.predict(features[left_out_indices])  # pyrefly: ignore[unbound-name]
   assert not np.any(np.isnan(scores))
 
   # Generate ROC curves with scores.
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/plotting.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/plotting.py
index 275de024..b6e21922 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/plotting.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/plotting.py
@@ -65,8 +65,8 @@ def plot_histograms(train: Iterable[float],
                     xlabel: Text = 'x',
                     thresh: Optional[float] = None) -> plt.Figure:
   """Plot histograms of training versus test metrics."""
-  xmin = min(np.min(train), np.min(test))
-  xmax = max(np.max(train), np.max(test))
+  xmin = min(np.min(train), np.min(test))  # pyrefly: ignore[no-matching-overload]
+  xmax = max(np.max(train), np.max(test))  # pyrefly: ignore[no-matching-overload]
   bins = np.linspace(xmin, xmax, 100)
   fig = plt.figure()
   plt.hist(test, bins=bins, density=True, alpha=0.5, label='test', log='y')
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/seq2seq_mia.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/seq2seq_mia.py
index 6f7fc82a..3512b049 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/seq2seq_mia.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/seq2seq_mia.py
@@ -277,9 +277,9 @@ def run_seq2seq_attack(attack_input: Seq2SeqAttackInputData,
   """
   attack_input.validate()
   attack_input_train, loss_train, accuracy_train = _get_attack_features_and_metadata(
-      attack_input.logits_train, attack_input.labels_train)
+      attack_input.logits_train, attack_input.labels_train)  # pyrefly: ignore[bad-argument-type]
   attack_input_test, loss_test, accuracy_test = _get_attack_features_and_metadata(
-      attack_input.logits_test, attack_input.labels_test)
+      attack_input.logits_test, attack_input.labels_test)  # pyrefly: ignore[bad-argument-type]
 
   privacy_report_metadata = privacy_report_metadata or PrivacyReportMetadata()
   privacy_report_metadata.loss_train = loss_train
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/tf_estimator_evaluation_example.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/tf_estimator_evaluation_example.py
index b3b4ffa9..7f0feaca 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/tf_estimator_evaluation_example.py
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/tf_estimator_evaluation_example.py
@@ -44,9 +44,9 @@ def small_cnn_fn(features, labels, mode):
     y = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(input_layer)
     y = tf.keras.layers.MaxPool2D()(y)
 
-  y = tf.keras.layers.Flatten()(y)
-  y = tf.keras.layers.Dense(64, activation='relu')(y)
-  logits = tf.keras.layers.Dense(10)(y)
+  y = tf.keras.layers.Flatten()(y)  # pyrefly: ignore[not-callable]
+  y = tf.keras.layers.Dense(64, activation='relu')(y)  # pyrefly: ignore[not-callable]
+  logits = tf.keras.layers.Dense(10)(y)  # pyrefly: ignore[not-callable]
 
   if mode != tf_estimator.ModeKeys.PREDICT:
     vector_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
@@ -58,7 +58,7 @@ def small_cnn_fn(features, labels, mode):
     optimizer = tf.train.MomentumOptimizer(
         learning_rate=FLAGS.learning_rate, momentum=0.9)
     global_step = tf.compat.v1.train.get_global_step()
-    train_op = optimizer.minimize(loss=scalar_loss, global_step=global_step)
+    train_op = optimizer.minimize(loss=scalar_loss, global_step=global_step)  # pyrefly: ignore[unbound-name]
     return tf_estimator.EstimatorSpec(
         mode=mode, loss=scalar_loss, train_op=train_op)
 
@@ -70,7 +70,7 @@ def small_cnn_fn(features, labels, mode):
                 labels=labels, predictions=tf.argmax(input=logits, axis=1))
     }
     return tf_estimator.EstimatorSpec(
-        mode=mode, loss=scalar_loss, eval_metric_ops=eval_metric_ops)
+        mode=mode, loss=scalar_loss, eval_metric_ops=eval_metric_ops)  # pyrefly: ignore[unbound-name]
 
   # Output the prediction probability (for PREDICT mode).
   elif mode == tf_estimator.ModeKeys.PREDICT:
diff --git a/tensorflow_privacy/privacy/privacy_tests/utils.py b/tensorflow_privacy/privacy/privacy_tests/utils.py
index 607500ab..5e3577e3 100644
--- a/tensorflow_privacy/privacy/privacy_tests/utils.py
+++ b/tensorflow_privacy/privacy/privacy_tests/utils.py
@@ -65,7 +65,7 @@ def log_loss(labels: np.ndarray,
   classes = np.unique(labels)
   if sample_weight is None:
     # If sample weights are not provided, set them to 1.0.
-    sample_weight = 1.0
+    sample_weight = 1.0  # pyrefly: ignore[bad-assignment]
   else:
     if np.shape(sample_weight)[0] != np.shape(labels)[0]:
       # Number of elements should be the same.
@@ -128,7 +128,7 @@ def squared_loss(y_true: np.ndarray,
                      'are %s and %s.' % (y_true.shape, y_pred.shape))
   if sample_weight is None:
     # If sample weights are not provided, set them to 1.0.
-    sample_weight = 1.0
+    sample_weight = 1.0  # pyrefly: ignore[bad-assignment]
   return sample_weight * (y_true - y_pred)**2
 
 
@@ -175,7 +175,7 @@ def multilabel_bce_loss(labels: np.ndarray,
                       '`from_logits` is set to False.'))
   if sample_weight is None:
     # If sample weights are not provided, set them to 1.0.
-    sample_weight = 1.0
+    sample_weight = 1.0  # pyrefly: ignore[bad-assignment]
   if isinstance(sample_weight, list):
     sample_weight = np.asarray(sample_weight)
   if isinstance(sample_weight, np.ndarray) and (sample_weight.ndim == 1):
@@ -245,15 +245,15 @@ def get_loss(
     loss_function = string_to_loss_function(loss_function)
   if loss_function == LossFunction.CROSS_ENTROPY:
     if multilabel_data:
-      loss = multilabel_bce_loss(labels, predictions, sample_weight,
+      loss = multilabel_bce_loss(labels, predictions, sample_weight,  # pyrefly: ignore[bad-argument-type]
                                  loss_function_using_logits)
     else:
-      loss = log_loss(labels, predictions, sample_weight,
+      loss = log_loss(labels, predictions, sample_weight,  # pyrefly: ignore[bad-argument-type]
                       loss_function_using_logits)
   elif loss_function == LossFunction.SQUARED:
-    loss = squared_loss(labels, predictions, sample_weight)
+    loss = squared_loss(labels, predictions, sample_weight)  # pyrefly: ignore[bad-argument-type]
   else:
-    loss = loss_function(labels, predictions, sample_weight)
+    loss = loss_function(labels, predictions, sample_weight)  # pyrefly: ignore[bad-argument-type]
   return loss
 
 
diff --git a/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/embedding.py b/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/embedding.py
index fed47cd7..fb416d7d 100644
--- a/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/embedding.py
+++ b/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/embedding.py
@@ -64,7 +64,7 @@ def embedding_layer_contribution_histogram(
           "The experimental embedding feature "
           "'_use_one_hot_matmul' is not supported."
       )
-  input_ids = tf.squeeze(tf.cast(*input_args, tf.int32))
+  input_ids = tf.squeeze(tf.cast(*input_args, tf.int32))  # pyrefly: ignore[bad-argument-count]
 
   def count_contributions_fn(
       grad: type_aliases.SparseGradient,
diff --git a/tensorflow_privacy/privacy/sparsity_preserving_noise/sparse_noise_utils.py b/tensorflow_privacy/privacy/sparsity_preserving_noise/sparse_noise_utils.py
index 4b96b9a0..d4b0c2f3 100644
--- a/tensorflow_privacy/privacy/sparsity_preserving_noise/sparse_noise_utils.py
+++ b/tensorflow_privacy/privacy/sparsity_preserving_noise/sparse_noise_utils.py
@@ -98,7 +98,7 @@ def _sample_sparse_indices_batch_size_heuristic(
   Returns:
     The batch size to use for sampling.
   """
-  max_num_samples = tf.cast(max_index + 1, tf.float32)
+  max_num_samples = tf.cast(max_index + 1, tf.float32)  # pyrefly: ignore[unsupported-operation]
   expected_num_samples = max_num_samples * probability
   # For expected samples > 50, choosing a batch size of 1.2 * expected samples
   # will allow for sampling only once to get all indices >95% of the time.
@@ -151,7 +151,7 @@ def sample_false_positive_indices(
 
   sampled_indices = tf.TensorArray(tf.int64, size=0, dynamic_size=True)
 
-  batch_size = batch_size or _sample_sparse_indices_batch_size_heuristic(
+  batch_size = batch_size or _sample_sparse_indices_batch_size_heuristic(  # pyrefly: ignore[bad-assignment]
       max_index, probability
   )
 
@@ -341,7 +341,7 @@ def extract_varname_to_contribution_counts_fns(
   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]
 
   varname_to_contribution_counts_fns = collections.defaultdict(list)
   for registry_fn_output in registry_fn_outputs_list:
diff --git a/tutorials/lm_dpsgd_tutorial.py b/tutorials/lm_dpsgd_tutorial.py
index e3302b87..3331c0a2 100644
--- a/tutorials/lm_dpsgd_tutorial.py
+++ b/tutorials/lm_dpsgd_tutorial.py
@@ -73,8 +73,8 @@ def rnn_model_fn(features, labels, mode):  # pylint: disable=unused-argument
   x = tf.reshape(x, [-1, SEQ_LEN])
   input_layer = x[:, :-1]
   input_one_hot = tf.one_hot(input_layer, 256)
-  lstm = tf.keras.layers.LSTM(256, return_sequences=True)(input_one_hot)
-  logits = tf.keras.layers.Dense(256)(lstm)
+  lstm = tf.keras.layers.LSTM(256, return_sequences=True)(input_one_hot)  # pyrefly: ignore[not-callable]
+  logits = tf.keras.layers.Dense(256)(lstm)  # pyrefly: ignore[not-callable]
 
   # Calculate loss as a vector (to support microbatches in DP-SGD).
   vector_loss = tf.nn.softmax_cross_entropy_with_logits(
@@ -122,11 +122,11 @@ def load_data():
           'using a substitute dataset from the tensorflow_datasets module.')
     train_dataset = tfds.load(
         name='lm1b/subwords8k',
-        split=tfds.Split.TRAIN,
+        split=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
         batch_size=NB_TRAIN,
         shuffle_files=True)
     test_dataset = tfds.load(
-        name='lm1b/subwords8k', split=tfds.Split.TEST, batch_size=10000)
+        name='lm1b/subwords8k', split=tfds.Split.TEST, batch_size=10000)  # pyrefly: ignore[missing-attribute]
     train_data = next(iter(tfds.as_numpy(train_dataset)))
     test_data = next(iter(tfds.as_numpy(test_dataset)))
     train_data = train_data['text'].flatten()
diff --git a/tutorials/mnist_dpsgd_tutorial_common.py b/tutorials/mnist_dpsgd_tutorial_common.py
index dac0cb23..20560d5f 100644
--- a/tutorials/mnist_dpsgd_tutorial_common.py
+++ b/tutorials/mnist_dpsgd_tutorial_common.py
@@ -28,9 +28,9 @@ def get_cnn_model(features):
       32, 4, strides=2, padding='valid', activation='relu')(
           y)
   y = tf.keras.layers.MaxPool2D(2, 1)(y)
-  y = tf.keras.layers.Flatten()(y)
-  y = tf.keras.layers.Dense(32, activation='relu')(y)
-  logits = tf.keras.layers.Dense(10)(y)
+  y = tf.keras.layers.Flatten()(y)  # pyrefly: ignore[not-callable]
+  y = tf.keras.layers.Dense(32, activation='relu')(y)  # pyrefly: ignore[not-callable]
+  logits = tf.keras.layers.Dense(10)(y)  # pyrefly: ignore[not-callable]
 
   return logits
 
@@ -40,7 +40,7 @@ def make_input_fn(split, input_batch_size=256, repetitions=-1, tpu=False):
 
   def input_fn(params=None):
     """A simple input function."""
-    batch_size = params.get('batch_size', input_batch_size)
+    batch_size = params.get('batch_size', input_batch_size)  # pyrefly: ignore[missing-attribute]
 
     def parser(example):
       image, label = example['image'], example['label']
diff --git a/tutorials/mnist_dpsgd_tutorial_eager.py b/tutorials/mnist_dpsgd_tutorial_eager.py
index 6948f882..f136c7b0 100644
--- a/tutorials/mnist_dpsgd_tutorial_eager.py
+++ b/tutorials/mnist_dpsgd_tutorial_eager.py
@@ -69,13 +69,13 @@ def main(_):
 
   # Create a dataset object and batch for the training data
   dataset = tf.data.Dataset.from_tensor_slices(
-      (tf.cast(train_images[..., tf.newaxis] / 255,
+      (tf.cast(train_images[..., tf.newaxis] / 255,  # pyrefly: ignore[bad-argument-type]
                tf.float32), tf.cast(train_labels, tf.int64)))
   dataset = dataset.shuffle(1000).batch(FLAGS.batch_size)
 
   # Create a dataset object and batch for the test data
   eval_dataset = tf.data.Dataset.from_tensor_slices(
-      (tf.cast(test_images[..., tf.newaxis] / 255,
+      (tf.cast(test_images[..., tf.newaxis] / 255,  # pyrefly: ignore[bad-argument-type]
                tf.float32), tf.cast(test_labels, tf.int64)))
   eval_dataset = eval_dataset.batch(10000)
 
@@ -133,7 +133,7 @@ def loss_fn():
     for (_, (images, labels)) in enumerate(eval_dataset.take(-1)):
       logits = mnist_model(images, training=False)
       correct_preds = tf.equal(tf.argmax(input=logits, axis=1), labels)
-    test_accuracy = np.mean(correct_preds.numpy())
+    test_accuracy = np.mean(correct_preds.numpy())  # pyrefly: ignore[unbound-name]
     print('Test accuracy after epoch %d is: %.3f' % (epoch, test_accuracy))
 
     # Compute the privacy budget expended so far.
diff --git a/tutorials/mnist_dpsgd_tutorial_vectorized.py b/tutorials/mnist_dpsgd_tutorial_vectorized.py
index e02bff5d..b0e84c9b 100644
--- a/tutorials/mnist_dpsgd_tutorial_vectorized.py
+++ b/tutorials/mnist_dpsgd_tutorial_vectorized.py
@@ -77,9 +77,9 @@ def cnn_model_fn(features, labels, mode):
       32, 4, strides=2, padding='valid', activation='relu')(
           y)
   y = tf.keras.layers.MaxPool2D(2, 1)(y)
-  y = tf.keras.layers.Flatten()(y)
-  y = tf.keras.layers.Dense(32, activation='relu')(y)
-  logits = tf.keras.layers.Dense(10)(y)
+  y = tf.keras.layers.Flatten()(y)  # pyrefly: ignore[not-callable]
+  y = tf.keras.layers.Dense(32, activation='relu')(y)  # pyrefly: ignore[not-callable]
+  logits = tf.keras.layers.Dense(10)(y)  # pyrefly: ignore[not-callable]
 
   # Calculate loss as a vector (to support microbatches in DP-SGD).
   vector_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
diff --git a/tutorials/mnist_lr_tutorial.py b/tutorials/mnist_lr_tutorial.py
index fce9ba76..ed6bba08 100644
--- a/tutorials/mnist_lr_tutorial.py
+++ b/tutorials/mnist_lr_tutorial.py
@@ -53,7 +53,7 @@
 def lr_model_fn(features, labels, mode, nclasses, dim):
   """Model function for logistic regression."""
   input_layer = tf.reshape(features['x'], tuple([-1]) + dim)
-  logits = tf.keras.layers.Dense(
+  logits = tf.keras.layers.Dense(  # pyrefly: ignore[not-callable]
       units=nclasses,
       kernel_regularizer=tf.keras.regularizers.L2(l2=FLAGS.regularizer),
       bias_regularizer=tf.keras.regularizers.L2(l2=FLAGS.regularizer))(
@@ -165,10 +165,10 @@ def print_privacy_guarantees(epochs, batch_size, samples, noise_multiplier):
     # Using RDP accountant to compute eps. Doing computation analytically is
     # an option.
     rdp = [order * coef for order in orders]
-    eps, _ = dp_accounting.rdp.compute_epsilon(orders, rdp, delta)
+    eps, _ = dp_accounting.rdp.compute_epsilon(orders, rdp, delta)  # pyrefly: ignore[bad-argument-type]
     print('\t{:g}% enjoy at least ({:.2f}, {})-DP'.format(p * 100, eps, delta))
 
-  accountant = dp_accounting.rdp.RdpAccountant(orders)
+  accountant = dp_accounting.rdp.RdpAccountant(orders)  # pyrefly: ignore[bad-argument-type]
   event = dp_accounting.SelfComposedDpEvent(
       dp_accounting.PoissonSampledDpEvent(
           batch_size / samples,
diff --git a/tutorials/movielens_tutorial.py b/tutorials/movielens_tutorial.py
index f39ec50d..af6f49bf 100644
--- a/tutorials/movielens_tutorial.py
+++ b/tutorials/movielens_tutorial.py
@@ -66,21 +66,21 @@ def nn_model_fn(features, labels, mode):
 
   # GMF part
   # Flatten the embedding vector as latent features in GMF
-  mf_user_latent = tf.keras.layers.Flatten()(mf_embedding_user(user_input))
-  mf_item_latent = tf.keras.layers.Flatten()(mf_embedding_item(item_input))
+  mf_user_latent = tf.keras.layers.Flatten()(mf_embedding_user(user_input))  # pyrefly: ignore[not-callable]
+  mf_item_latent = tf.keras.layers.Flatten()(mf_embedding_item(item_input))  # pyrefly: ignore[not-callable]
   # Element-wise multiply
   mf_vector = tf.keras.layers.multiply([mf_user_latent, mf_item_latent])
 
   # MLP part
   # Flatten the embedding vector as latent features in MLP
-  mlp_user_latent = tf.keras.layers.Flatten()(mlp_embedding_user(user_input))
-  mlp_item_latent = tf.keras.layers.Flatten()(mlp_embedding_item(item_input))
+  mlp_user_latent = tf.keras.layers.Flatten()(mlp_embedding_user(user_input))  # pyrefly: ignore[not-callable]
+  mlp_item_latent = tf.keras.layers.Flatten()(mlp_embedding_item(item_input))  # pyrefly: ignore[not-callable]
   # Concatenation of two latent features
   mlp_vector = tf.keras.layers.concatenate([mlp_user_latent, mlp_item_latent])
 
   predict_vector = tf.keras.layers.concatenate([mf_vector, mlp_vector])
 
-  logits = tf.keras.layers.Dense(5)(predict_vector)
+  logits = tf.keras.layers.Dense(5)(predict_vector)  # pyrefly: ignore[not-callable]
 
   # Calculate loss as a vector (to support microbatches in DP-SGD).
   vector_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(