Update Keras layers to use `tf.xyz` rather than `tf.raw_ops.xyz`. (E.g., use...

Update Keras layers to use `tf.xyz` rather than `tf.raw_ops.xyz`.  (E.g., use tf.matmul rather than tf.raw_ops.MatMul.)

ExtensionTypes that want to support these layers can do so by adding dispatch.  (E.g., tf.matmul supports dispatch; but tf.raw_ops.MatMul does not.)

PiperOrigin-RevId: 396441290

keras/layers/core/dense.py

@@ -196,7 +196,7 @@ class Dense(Layer):
         outputs = tf.nn.embedding_lookup_sparse(
             self.kernel, ids, weights, combiner='sum')
       else:
-        outputs = tf.raw_ops.MatMul(a=inputs, b=self.kernel)
+        outputs = tf.matmul(a=inputs, b=self.kernel)
     # Broadcast kernel to inputs.
     else:
       outputs = tf.tensordot(inputs, self.kernel, [[rank - 1], [0]])

keras/layers/kernelized.py

@@ -210,7 +210,7 @@ class RandomFourierFeatures(base_layer.Layer):
     inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
     inputs = tf.cast(inputs, tf.float32)
     kernel = (1.0 / self.kernel_scale) * self.unscaled_kernel
-    outputs = tf.raw_ops.MatMul(a=inputs, b=kernel)
+    outputs = tf.matmul(a=inputs, b=kernel)
     outputs = tf.nn.bias_add(outputs, self.bias)
     return tf.cos(outputs)
 

keras/layers/local.py

@@ -775,10 +775,8 @@ def local_conv_sparse_matmul(inputs, kernel, kernel_idxs, kernel_shape,
       Output (N+2)-D dense tensor with shape `output_shape`.
   """
   inputs_flat = backend.reshape(inputs, (backend.shape(inputs)[0], -1))
-  output_flat = tf.raw_ops.SparseTensorDenseMatMul(
-      a_indices=kernel_idxs,
-      a_values=kernel,
-      a_shape=kernel_shape,
+  output_flat = tf.sparse.sparse_dense_matmul(
+      sp_a=tf.SparseTensor(kernel_idxs, kernel, kernel_shape),
       b=inputs_flat,
       adjoint_b=True)
   output_flat_transpose = backend.transpose(output_flat)

keras/utils/metrics_utils.py

@@ -642,7 +642,7 @@ def update_confusion_matrix_variables(variables_to_update,
   if y_pred.shape.ndims == 1:
     num_labels = 1
   else:
-    num_labels = tf.raw_ops.Prod(input=pred_shape[1:], axis=0)
+    num_labels = tf.math.reduce_prod(pred_shape[1:], axis=0)
   thresh_label_tile = tf.where(one_thresh, num_labels,
                                          tf.ones([], dtype=tf.int32))