code: def _slice_with_actions(embeddings, actions):
      """Slice a Tensor.
  
      Take embeddings of the form [batch_size, num_actions, embed_dim]
      and actions of the form [batch_size, 1], and return the sliced embeddings
      like embeddings[:, actions, :].
  
      Args:
        embeddings: Tensor of embeddings to index.
        actions: int Tensor to use as index into embeddings
  
      Returns:
        Tensor of embeddings indexed by actions
      """
      batch_size, num_actions = embeddings.get_shape()[:2]
  
      # Values are the 'values' in a sparse tensor we will be setting
      act_indx = tf.cast(actions, tf.int64)[:, None]
      values = tf.reshape(
          tf.cast(tf.ones(tf.shape(actions)), tf.bool), [-1]
      )
  
      # Create a range for each index into the batch
      act_range = tf.range(0, batch_size, dtype=tf.int64)[:, None]
      # Combine this into coordinates with the action indices
      indices = tf.concat([act_range, act_indx], 1)
  
      actions_mask = tf.SparseTensor(
          indices, values, [batch_size, num_actions]
      )
      actions_mask = tf.stop_gradient(
          tf.sparse_tensor_to_dense(actions_mask, default_value=False)
      )
      sliced_emb = tf.boolean_mask(embeddings, actions_mask)
      return sliced_emb
  
file_path: trfl/trfl/dist_value_ops.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import tensorflow.compat.v1
  from trfl import base_ops
  from trfl import distribution_ops
label_desc: Use image node_id 1 for calling a global function with example usage: _slice_with_actions(embeddings, actions) and returns: sliced_emb