code: def _index_embs_with_actions(
      embeddings: tf.Tensor,
      actions: tf.Tensor,
  ) -> tf.Tensor:
      """Slice an embedding Tensor with action indices.
  
      Take embeddings of the form [batch_size, num_actions, embed_dim]
      and actions of the form [batch_size], and return the sliced embeddings
      like embeddings[:, actions, :]. Doing this my way because the comments in
      the official op are scary.
  
      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.shape.as_list()
  
      # Values are the 'values' in a sparse tensor we will be setting
      act_indx = tf.cast(actions, tf.int64)[:, None]
      values = tf.ones([tf.size(actions)], dtype=tf.bool)
  
      # 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.to_dense(actions_mask, default_value=False)
      )
      sliced_emb = tf.boolean_mask(embeddings, actions_mask)
      return sliced_emb
  
imports: from typing import Dict, List, Optional, Tuple
  from acme import core
  from acme import types
  from acme.adders import reverb
  from acme.tf import losses
  from acme.tf import networks
  from acme.tf import savers
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import reverb
  import sonnet
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: _index_embs_with_actions(embeddings, actions) and returns: sliced_emb