def sp_sp_matmul(sp_a, sp_b):
  """Element-wise multiply two SparseTensors of same shape."""
  sp_a = tf.sparse.reorder(sp_a)
  sp_b = tf.sparse.reorder(sp_b)

  # 1. create 1d index maps
  scaling_vector = tf.cumprod(tf.cast(tf.shape(sp_b), sp_b.indices.dtype))
  a1s = tf.linalg.matvec(sp_a.indices, scaling_vector)
  b1s = tf.linalg.matvec(sp_b.indices, scaling_vector)

  # 2. get relevant indices of sp_a
  int_idx = tf.setdiff1d(a1s, b1s)[1]
  to_retain = tf.sparse_to_dense(
      sparse_indices=int_idx,
      output_shape=tf.shape(a1s),
      default_value=0.0,
      sparse_values=1.0) < 0.5
  rsp_a = tf.sparse.retain(sp_a, to_retain)
  rsp_a = tf.sparse.reorder(rsp_a)

  # 3. get relevant indices of sp_b
  int_idx = tf.setdiff1d(b1s, a1s)[1]
  to_retain = tf.sparse_to_dense(
      sparse_indices=int_idx,
      output_shape=tf.shape(b1s),
      default_value=0.0,
      sparse_values=1.0) < 0.5
  rsp_b = tf.sparse.retain(sp_b, to_retain)
  rsp_b = tf.sparse.reorder(rsp_b)

  # 4. create output matrix
  return tf.SparseTensor(
      indices=rsp_b.indices,
      values=rsp_a.values * rsp_b.values,
      dense_shape=sp_b.dense_shape)
Element-wise multiply two SparseTensors of same shape.
language/labs/drkit/model_fns.py