def Take(*args):  # pylint: disable=invalid-name
  """Layer to pick subset of inputs from parallel input stream.

  Args:
    *args: a sequence of ints

  Returns:
    A new layer that selects inputs from an incoming parallel stream.
    In numpy notation: outputs = parallel_inputs[args]
    If the resulting output list has only one member, it is automatically
    unwrapped and the contents are passed on directly.
  """
  def init_fun(_, input_shape):
    output_shape = []
    for arg in args:
      output_shape.append(input_shape[arg])
    if len(output_shape) == 1:
      output_shape = output_shape[0]
    return (output_shape, ())
  def apply_fun(params, inputs, **kwargs):
    del params, kwargs
    outputs = []
    for arg in args:
      outputs.append(inputs[arg])
    if len(outputs) == 1:
      outputs = outputs[0]
    return outputs
  return init_fun, apply_fun