code: def _update_timestep(x, timestep, values):
      """Set x[:, timestep] = values.
  
      This operation is **NOT** differentiable.
  
      Args:
        x: Tensor of shape [batch_size, seq_len, ...]
        timestep: int or scalar Tensor. Index to update in x.
        values: Tensor of shape [batch_size, ...]. New values for x[:, i].
  
      Returns:
        Copy of 'x' after setting x[:, timestep] = values.
      """
      perm = range(x.shape.ndims)
      perm[0], perm[1] = perm[1], perm[0]
      x = tf.transpose(x, perm)
      x = inplace_ops.alias_inplace_update(x, timestep, values)
      x = tf.transpose(x, perm)
      return x
  
file_path: tensor2tensor/tensor2tensor/utils/scheduled_sampling.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.ops import inplace_ops
label_desc: Use image node_id 5 for calling a global function with example usage: _update_timestep(x, timestep, values) and returns: x