def __init__(
      self, body_value, body_start=None, tail_start=None, tail_fn=None):
    """Specifies a body-and-tail time curve.

    Args:
      body_value: Constant learning rate for the body of the curve (after
          warm-up and before tail). Also is the reference (maximum) value for
          calculating warm-up values and tail values.
      body_start: Training step number at which the body starts. If None, takes
          its value from tail_start, which amounts to there being no body. All
          steps from 1 to body_start - 1 are computed using a linear warm-up.
      tail_start: Training step number at which the tail starts. If None, the
          body value remains until the end of training.
      tail_fn: Function returning a floating point learning rate, given inputs:
            - step_number (absolute step number from the start of training)
            - tail_start (step number at which the tail starts)
            - body_value (value relative to which the tail should be computed)
    """
    if body_start is None and tail_start is None:
      raise ValueError('Both body start and tail start are None.')
    if tail_start is not None and tail_fn is None:
      raise ValueError(
          f'Tail start has value ({tail_start}) but tail_fn is None.')
    if body_start is None:
      body_start = tail_start if tail_start is not None else 1

    self._body_value = body_value
    self._body_start = body_start
    self._tail_start = tail_start
    self._tail_fn = tail_fn
Specifies a body-and-tail time curve.

    Args:
      body_value: Constant learning rate for the body of the curve (after
          warm-up and before tail). Also is the reference (maximum) value for
          calculating warm-up values and tail values.
      body_start: Training step number at which the body starts. If None, takes
          its value from tail_start, which amounts to there being no body. All
          steps from 1 to body_start - 1 are computed using a linear warm-up.
      tail_start: Training step number at which the tail starts. If None, the
          body value remains until the end of training.
      tail_fn: Function returning a floating point learning rate, given inputs:
            - step_number (absolute step number from the start of training)
            - tail_start (step number at which the tail starts)
            - body_value (value relative to which the tail should be computed)
    
trax/supervised/lr_schedules.py