def __init__(self, axes, max_lengths=None, **kwargs):
    """Init.

    Args:
      axes: list of ints, axis over which to apply the positional embeddings.
      max_lengths: list of ints, maximum length over each axis.
      **kwargs:
    """
    super(PositionEmbed, self).__init__(**kwargs)
    if not isinstance(axes, (list, tuple)):
      axes = [axes]
    self.axes = axes
    self.max_lengths = None
    if max_lengths:
      if not isinstance(max_lengths, (list, tuple)):
        max_lengths = [max_lengths]
      self.max_lengths = max_lengths