def layer(output_shape=None, new_parameters=None):
  """Create a layer class from a function."""
  def layer_decorator(call):
    """Decorating the call function."""
    def output_shape_fun(self, input_shape):
      if output_shape is None:
        return input_shape
      kwargs = self._init_kwargs  # pylint: disable=protected-access
      return output_shape(input_shape, **kwargs)

    def new_parameters_fun(self, input_shape, rng):
      if new_parameters is None:
        return ()
      kwargs = self._init_kwargs  # pylint: disable=protected-access
      return new_parameters(input_shape, rng, **kwargs)

    def call_fun(self, x, params=(), **kwargs):
      """The call function of the created class, derived from call."""
      # Merge on-call kwargs with class-kwargs.
      call_kwargs = kwargs.copy()
      call_kwargs.update(self._init_kwargs)  # pylint: disable=protected-access
      # Call with the merged kwargs.
      return call(x, params=params, **call_kwargs)

    # Set doc for python help.
    call_fun.__doc__ = call.__doc__
    if output_shape is None:
      output_shape_fun.__doc__ = output_shape.__doc__
    if new_parameters is None:
      new_parameters_fun.__doc__ = new_parameters.__doc__

    # Create the class.
    cls = type(call.__name__, (Layer,),
               {'call': call_fun,
                'output_shape': output_shape_fun,
                'new_parameters': new_parameters_fun})

    return cls
  return layer_decorator
Create a layer class from a function.