def remove_from_sparse(sp_tensor, remove_indices):
  """Remove indices from SparseTensor if present."""
  # 1. create 1d index maps
  scaling_vector = tf.cumprod(
      tf.cast(tf.shape(sp_tensor), sp_tensor.indices.dtype))
  a1s = tf.linalg.matvec(sp_tensor.indices, scaling_vector)
  b1s = tf.linalg.matvec(remove_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
  return tf.sparse.retain(sp_tensor, to_retain)
Remove indices from SparseTensor if present.
language/labs/drkit/model_fns.py