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

  Args:
    d_feature: Number of memory channels (dimensionality of feature embedding).
    steps: Number of times depthwise recurrence steps.
    vocab_size: Vocabulary size.
    mode: Whether we are training or evaluating or doing inference.

  Returns:
    A NeuralGPU Stax model.
  """
  del mode

  core = ConvDiagonalGRU(units=d_feature)
  return tl.Serial(
      tl.Embedding(vocab_size=vocab_size, d_feature=d_feature),
      [core] * steps,
      tl.Dense(vocab_size),
  )
Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.

  Args:
    d_feature: Number of memory channels (dimensionality of feature embedding).
    steps: Number of times depthwise recurrence steps.
    vocab_size: Vocabulary size.
    mode: Whether we are training or evaluating or doing inference.

  Returns:
    A NeuralGPU Stax model.
  
trax/models/neural_gpu.py