def create_cm_sketch(topk_obj_ids,
                     topk_obj_weights,
                     all_entity_sketches,
                     cm_width):
  """Create cm sketches for a set of weighted entities.

  Args:
    topk_obj_ids: batch_size, topk
    topk_obj_weights: batch_size, topk
    all_entity_sketches: num_entities, depth
    cm_width: width of count-min sketch

  Returns:
    k hot dense vectors: batch_size, depth, width
  """
  topk_fact_obj_sketches = tf.gather(
      all_entity_sketches, topk_obj_ids, axis=0)
  # batch_size, topk, depth
  batch_size = tf.shape(topk_fact_obj_sketches)[0]
  topk = tf.shape(topk_fact_obj_sketches)[1]
  cm_depth = tf.shape(topk_fact_obj_sketches)[2]

  # We first create a sparse matrix from the hash values. We will then
  # convert it into dense matrix. This is more efficient than creating
  # k one-hot vectors and then aggregating them into one k-hot vector.

  # First prepare ids of non-zero values in the sparse matrix
  flattened_topk_hash_ids = tf.reshape(topk_fact_obj_sketches, shape=[-1])
  # batch_size * topk * depth
  topk_obj_weights = tf.tile(
      tf.expand_dims(topk_obj_weights, axis=2),
      multiples=[1, 1, cm_depth])
  # batch_size, topk, depth
  flattened_topk_obj_weights = tf.reshape(
      topk_obj_weights, shape=[-1])
  # batch_size * topk * depth
  batch_ids = tf.range(batch_size)
  # batch_size,
  batch_ids = tf.expand_dims(tf.expand_dims(batch_ids, axis=1), axis=2)
  # batch_size, 1, 1
  batch_ids = tf.tile(batch_ids, multiples=[1, topk, cm_depth])
  # batch_size, topk, depth
  flattened_batch_ids = tf.reshape(batch_ids, shape=[-1])
  # batch_size * topk * depth
  depth_ids = tf.range(cm_depth)
  # depth,
  depth_ids = tf.expand_dims(tf.expand_dims(depth_ids, axis=0), axis=1)
  # 1, 1, depth
  depth_ids = tf.tile(depth_ids, multiples=[batch_size, topk, 1])
  # batch_size, topk, depth
  flattened_depth_ids = tf.reshape(depth_ids, shape=[-1])
  # batch_size * topk * depth
  sparse_value_ids = tf.cast(tf.stack(
      [flattened_batch_ids, flattened_depth_ids, flattened_topk_hash_ids],
      axis=1), dtype=tf.int64)

  # Then prepare values of non-zero values in the sparse matrix. Values
  # are sorted to ascending order. If there are duplicates, later (larger)
  # values will be kept.
  sorted_orders = tf.argsort(
      flattened_topk_obj_weights, direction='ASCENDING', stable=True)
  # batch_size * topk * depth
  sorted_flattened_topk_obj_weights = tf.gather(
      flattened_topk_obj_weights, sorted_orders)
  sorted_sparse_value_ids = tf.gather(sparse_value_ids, sorted_orders)

  # Finally create sketch in sparse tensors and convert it to dense tensors.
  # We donot validate indices here. If multiple values are about to be assigned
  # to the same row and column, we will keep the last value, because the last
  # value is the larger one. This behaviour is by design.
  sparse_k_hot_sketch = tf.SparseTensor(
      indices=sorted_sparse_value_ids,
      values=sorted_flattened_topk_obj_weights,
      dense_shape=[batch_size, cm_depth, cm_width])
  dense_k_hot_sketch = tf.sparse.to_dense(
      sparse_k_hot_sketch, validate_indices=False)
  # batch_size, cm_depth, cm_width
  return dense_k_hot_sketch
Create cm sketches for a set of weighted entities.

  Args:
    topk_obj_ids: batch_size, topk
    topk_obj_weights: batch_size, topk
    all_entity_sketches: num_entities, depth
    cm_width: width of count-min sketch

  Returns:
    k hot dense vectors: batch_size, depth, width
  
language/emql/module.py