def register_recurrent_cell(state_fn, **kwargs):
  """Decorator to register the current cell function with auxiliary functions.

  The only current auxauxiliaryilary function is `sequence_CELL_NAME` which
  applies the cell function across a sequence.

  Note, the cell_function must have arguments like:

      cell_fn(input_, state_list, *args, **kwargs)

  If it contains a field called 'lengths', then it is used to intelligently stop
  long unrolls.

  It will be registered as a compound op, please see `RegisterCompoundOp` for
  details.

  state_fn will then be called as:

      state_fn(input_, *args, **kwargs)

  Args:
    state_fn: An function that returns a list of state tuples
        `[(name, size_ex_batch)]`.
    **kwargs: Arguments to pass forward to the registration methods.

  Returns:
    The actual decorator, which in turn returns the registered cell function.
  """
  def _fn(cell_fn):
    """The register_recurrent_cell decorator."""
    result = prettytensor.RegisterCompoundOp(**kwargs)(cell_fn)

    method_name = kwargs.pop('method_name', cell_fn.__name__)
    if method_name.endswith('_cell'):
      method_name = method_name[:-5]
    method_name = 'sequence_' + method_name

    class _Sequence(prettytensor.VarStoreMethod):
      """Expands a cell into a sequence version.

      This holds an internal template to ensure that the same parameters are
      used each time it is called.
      """

      def __init__(self):
        super(self.__class__, self).__init__()
        self.template = None

      def __call__(self,
                   input_layer,
                   *args,
                   **kwargs):
        name = kwargs.pop('name', 'sequence')
        lengths = kwargs.pop('lengths', None)
        state_tuples = state_fn(input_layer, *args, **kwargs)
        if not self.template:
          template = prettytensor.template('input', input_layer.bookkeeper)
          states = [prettytensor.UnboundVariable(s[0]) for s in state_tuples]
          self.template = cell_fn(template, states, *args, **kwargs).as_fn(
              'input', *[s[0] for s in state_tuples])

        batch_size = input_layer.shape[0]
        state_shapes = [[batch_size, s[1]] for s in state_tuples]
        return unroll_state_saver(
            input_layer, name, state_shapes, self.template, lengths)

    doc_to_copy = result.__doc__
    doc_to_copy = doc_to_copy[doc_to_copy.find('\n') + 2:]
    doc = """Unrolls `{0}` over the input.

  This takes an input that is a list of length timesteps where each element
  is a `Tensor` of `[batch, *Dims]` and unrolls the recurrent cell. The input
  and state to the cell are managed by this method, but the rest of the
  arguments are passed through.

    {1}""".format(cell_fn.__name__, doc_to_copy)
    if hasattr(_Sequence.__call__, '__func__'):
      _Sequence.__call__.__func__.__doc__ = doc
    else:
      _Sequence.__call__.__doc__ = doc
    prettytensor.Register(method_name=method_name, **kwargs)(_Sequence)

    return result
  return _fn
Decorator to register the current cell function with auxiliary functions.

  The only current auxauxiliaryilary function is `sequence_CELL_NAME` which
  applies the cell function across a sequence.

  Note, the cell_function must have arguments like:

      cell_fn(input_, state_list, *args, **kwargs)

  If it contains a field called 'lengths', then it is used to intelligently stop
  long unrolls.

  It will be registered as a compound op, please see `RegisterCompoundOp` for
  details.

  state_fn will then be called as:

      state_fn(input_, *args, **kwargs)

  Args:
    state_fn: An function that returns a list of state tuples
        `[(name, size_ex_batch)]`.
    **kwargs: Arguments to pass forward to the registration methods.

  Returns:
    The actual decorator, which in turn returns the registered cell function.