Resolve type mismatching bugs between given model weights and pruning variables of pruning api

PiperOrigin-RevId: 414446051

Author: rino20

tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py

@@ -98,15 +98,30 @@ def _update_mask(self, weights):
                   (1 - sparsity)),
               1),
           tf.int32)
-      # Sort the entire array
-      values, _ = tf.math.top_k(
+      values, indices = tf.math.top_k(
           tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))
-      # Grab the (k-1)th value
 
-      current_threshold = tf.gather(values, k - 1)
-      new_mask = tf.dtypes.cast(
-          tf.math.greater_equal(abs_weights, current_threshold), weights.dtype)
-    return current_threshold, new_mask
+      # Grab the (k-1)th value as a threshold to build pruning mask.
+      threshold_value = tf.gather(values, k - 1)
+      threshold_pos = tf.gather(indices, k - 1)
+
+      # Build mask for the weight element higher magnitude than threshold value.
+      # A mask is added to make sure the threshold element be incorporated.
+      # TODO(b/208967539): Update the logic to index oriented logic.
+      new_mask = tf.math.logical_or(
+          tf.math.greater_equal(abs_weights, threshold_value),
+          tf.reshape(
+              tf.one_hot(
+                  threshold_pos,
+                  depth=tf.size(abs_weights),
+                  on_value=True,
+                  off_value=False,
+                  dtype=tf.bool), abs_weights.shape))
+
+    # Updated mask is casted back to weight's data type in case of the type
+    # mismatching due to keras mixed precision policy.
+    return tf.dtypes.cast(threshold_value, weights.dtype), tf.dtypes.cast(
+        new_mask, weights.dtype)
 
   def _update_mask_sparsity_m_by_n(self, weights, m_by_n=(2, 4)):
     """Updates the m by n sparsity mask for a given weight tensor.
@@ -221,14 +236,16 @@ def update_var(variable, reduced_value):
     if tf.distribute.get_replica_context():
       values_and_vars = []
       for weight, mask, _ in self._pruning_vars:
-        masked_weight = tf.math.multiply(weight, mask)
+        masked_weight = tf.dtypes.cast(
+            tf.math.multiply(weight, mask), dtype=weight.dtype)
         values_and_vars.append((masked_weight, weight))
       if values_and_vars:
         assign_objs.append(tf.distribute.get_replica_context().merge_call(
             update_fn, args=(values_and_vars,)))
     else:
       for weight, mask, _ in self._pruning_vars:
-        masked_weight = tf.math.multiply(weight, mask)
+        masked_weight = tf.dtypes.cast(
+            tf.math.multiply(weight, mask), dtype=weight.dtype)
         assign_objs.append(tf_compat.assign(weight, masked_weight))
 
     return assign_objs