code: def _update_row(input_x, updates, row_index):
      """Updates the i-th row of tensor `x` with the values given in `updates`.
  
      Args:
        input_x: the input tensor.
        updates: the values to place on the i-th row of `x`.
        row_index: which row to update.
  
      Returns:
        The updated tensor (same shape as `x`).
      """
      n = (
          tf.compat.dimension_value(input_x.shape[1])
          or tf.shape(input_x)[1]
      )
      indices = tf.concat(
          [
              row_index * tf.ones([n, 1], dtype=tf.int32),
              tf.reshape(tf.range(n, dtype=tf.int32), [n, 1]),
          ],
          axis=-1,
      )
      return tf.tensor_scatter_nd_update(
          tensor=input_x, indices=indices, updates=tf.squeeze(updates)
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import tensorflow
  from tf_agents.bandits.environments import non_stationary_stochastic_environment
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
label_desc: Use image node_id 2 for calling a global function with example usage: _update_row(input_x, updates, row_index) and returns: tf