Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.

  Args:
    feature_depth: Number of memory channels
    steps: Number of times depthwise recurrence steps.
    vocab_size: Vocabulary size.

  Returns:
    A NeuralGPU Stax model.

def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
  """Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.

  Args:
    feature_depth: Number of memory channels
    steps: Number of times depthwise recurrence steps.
    vocab_size: Vocabulary size.

  Returns:
    A NeuralGPU Stax model.
  """
  xs = []
  xs.append(
      layers.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
  core = ConvDiagonalGRU(units=feature_depth)
  xs.extend([core] * steps)
  xs.append(layers.Dense(vocab_size))
  xs.append(layers.LogSoftmax())

  return layers.Serial(*xs)