def sp_mul(sp_a_indices, sp_a_values, sp_b, default_value=0.):
  """Element-wise multiply two SparseTensors.

  One SparseTensor has a fixed number of non-zero values, represented in the
  form of a B x K matrix of indices and the other in a B x K matrix of values.

  Args:
    sp_a_indices: B x K Tensor. Indices into M values.
    sp_a_values: B x K Tensor.
    sp_b: B x M SparseTensor.
    default_value: Scalar default value to use for indices which do not
      intersect.

  Returns:
    mul_values: B x K Tensor corresponding to indices in sp_a_indices. For
      indices also present in sp_b the output is the product of the values in
      sp_a_values. For those absent in sp_b, the output is default_value.
  """
  batch_size = tf.shape(sp_a_indices)[0]
  num_neighbors = tf.shape(sp_a_indices)[1]
  a_values = tf.reshape(sp_a_values, (batch_size * num_neighbors,))
  brange = tf.tile(
      tf.expand_dims(tf.cast(tf.range(batch_size), sp_a_indices.dtype), 1),
      (1, num_neighbors))
  nnz_indices = tf.stack([brange, sp_a_indices], axis=2)
  nnz_indices = tf.reshape(nnz_indices, (batch_size * num_neighbors, 2))

  # Linearize
  scaling_vector = tf.cumprod(tf.cast(tf.shape(sp_b), sp_b.indices.dtype))
  a1s = tf.linalg.matvec(
      tf.cast(nnz_indices, sp_b.indices.dtype), scaling_vector)
  b1s = tf.linalg.matvec(sp_b.indices, scaling_vector)

  # Get intersections.
  a_nonint, _ = tf.setdiff1d(a1s, b1s)
  a_int, a_int_idx = tf.setdiff1d(a1s, a_nonint)
  b_nonint, _ = tf.setdiff1d(b1s, a1s)
  b_int, b_int_idx = tf.setdiff1d(b1s, b_nonint)

  # Get sorting.
  a_int_sort_ind = tf.argsort(a_int)
  b_int_sort_ind = tf.argsort(b_int)

  # Multiply.
  int_vals = (
      tf.gather(a_values, tf.gather(a_int_idx, a_int_sort_ind)) *
      tf.gather(sp_b.values, tf.gather(b_int_idx, b_int_sort_ind)))
  int_vals_dense = tf.sparse_to_dense(
      sparse_indices=tf.expand_dims(tf.gather(a_int_idx, a_int_sort_ind), 1),
      output_shape=tf.shape(a_values),
      default_value=default_value,
      sparse_values=int_vals,
      validate_indices=False)

  return tf.reshape(int_vals_dense, (batch_size, num_neighbors))
Element-wise multiply two SparseTensors.

  One SparseTensor has a fixed number of non-zero values, represented in the
  form of a B x K matrix of indices and the other in a B x K matrix of values.

  Args:
    sp_a_indices: B x K Tensor. Indices into M values.
    sp_a_values: B x K Tensor.
    sp_b: B x M SparseTensor.
    default_value: Scalar default value to use for indices which do not
      intersect.

  Returns:
    mul_values: B x K Tensor corresponding to indices in sp_a_indices. For
      indices also present in sp_b the output is the product of the values in
      sp_a_values. For those absent in sp_b, the output is default_value.
  
language/labs/drkit/model_fns.py