code: def _batched_table_lookup(tbl, row, col):
      """Mapped 2D table lookup.
  
      Args:
        tbl: a `Tensor` of shape `[r, s, t]`.
        row: a `Tensor` of dtype `int32` with shape `[r]` and values in the range
          `[0, s - 1]`.
        col: a `Tensor` of dtype `int32` with shape `[r]` and values in the range
          `[0, t - 1]`.
  
      Returns:
        A `Tensor` `x` with shape `[r]` where `x[i] = tbl[i, row[i], col[i]`.
      """
      assert_correct_shapes = tf.group(
          tf.assert_equal(tf.shape(row), tf.shape(col)),
          tf.assert_equal(tf.shape(row)[0], tf.shape(tbl)[0]),
      )
      rng = tf.range(tf.shape(row)[0])
      idx = tf.stack([rng, row, col], axis=-1)
      with tf.control_dependencies([assert_correct_shapes]):
          values = tf.gather_nd(tbl, idx)
      return values
  
file_path: agents/tf_agents/bandits/environments/classification_environment.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Text
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.environments import bandit_tf_environment
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import eager_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _batched_table_lookup(tbl, row, col) and returns: values