code: def Reorder(x, params, output=None, **kwargs):
    
    del params, kwargs
    if output is None:
      return x
    return base.nested_map(output, lambda i: x[i])
docs: Reorder a tuple into another tuple.
  
    For example, we can re-order (x, y) into (y, x) or even (y, (x, y), y).
    The output argument specifies how to re-order, using integers that refer
    to indices in the input tuple. For example, if
  
      input = (x, y, z)
  
    then
  
      Reorder(input, output=(1, 0, 2))   = (y, x, z)
      Reorder(input, output=(0, 0))      = (x, x)
      Reorder(input, output=(0, (1, 1))) = (x, (y, y))
      Reorder(input, output=((2, 0), (1, 1))) = ((z, x), (y, y))
  
    By default (if no output is given) Reorder does nothing (Identity).
  
    Args:
      x: the input tuple to re-order.
      params: layer parameters (unused).
      output: the specification of the output tuple: a nested tuple of ints.
      **kwargs: other arguments (unused).
  
    Returns:
      The re-ordered tuple with the same shape as output.