def get_item_at_decode_step(x, decode_step, query_shape):
  """Extracts a single item from an n-d tensor at `decode_step` position.

  Args:
    x: a [batch, d1, d2, ..., dn, depth] tensor
    decode_step: an integer
    query_shape: a tuple (q1, q2, ..., qn) representing the query shape

  Returns:
    a [batch, 1, 1, ..., 1, depth] tensor that is a single element from `x` at
    `decode_step` w.r.t. blocked raster scan order.
  """
  x_shape = common_layers.shape_list(x)
  index = decode_step_to_index(decode_step, query_shape, x_shape[1:-1])
  # TPU needs size to be non negative for the case when begins are not
  # compile-time constants.
  return tf.slice(x, [0] + list(index) + [0],
                  [x_shape[0]] + [1] * len(index) + [x_shape[-1]])
Extracts a single item from an n-d tensor at `decode_step` position.

  Args:
    x: a [batch, d1, d2, ..., dn, depth] tensor
    decode_step: an integer
    query_shape: a tuple (q1, q2, ..., qn) representing the query shape

  Returns:
    a [batch, 1, 1, ..., 1, depth] tensor that is a single element from `x` at
    `decode_step` w.r.t. blocked raster scan order.
  
tensor2tensor/layers/common_attention.py