def select_block_for_decode_step(blocked_x, decode_step, query_shape):
  """Selects one block from `x` that contains position `decode_step`.

  NOTE: This method only works for blocked inputs. It selects one block around
  `decode_step` position in blocked raster scan order.

  Args:
    blocked_x: a [batch, blocks_per_d1, ..., blocks_per_dn, b1 * ...* bn, depth]
      tensor
    decode_step: an integer
    query_shape: a tuple (q1, q2, ..., qn) representing query shape

  Returns:
     a [batch, [1] * n, b1 * ... * bn, depth] tensor
  """
  blocked_x_shape = common_layers.shape_list(blocked_x)
  # calculate the shape of the normal x
  x_shape = [b * q for b, q in zip(blocked_x_shape[1:-2], query_shape)]
  # Get the position of `decode_step` element in the unblocked x.
  index = decode_step_to_index(decode_step, query_shape, x_shape)
  # Convert it to the blocked positions.
  blocked_index = [i // q for i, q in zip(index, query_shape)]
  # TPU needs size to be non negative for the case when begin is not
  # compile-time constants.
  return tf.slice(blocked_x, [0] + blocked_index + [0, 0],
                  [blocked_x_shape[0]] + [1] * len(blocked_index) +
                  blocked_x_shape[-2:])
Selects one block from `x` that contains position `decode_step`.

  NOTE: This method only works for blocked inputs. It selects one block around
  `decode_step` position in blocked raster scan order.

  Args:
    blocked_x: a [batch, blocks_per_d1, ..., blocks_per_dn, b1 * ...* bn, depth]
      tensor
    decode_step: an integer
    query_shape: a tuple (q1, q2, ..., qn) representing query shape

  Returns:
     a [batch, [1] * n, b1 * ... * bn, depth] tensor
  
tensor2tensor/layers/common_attention.py