code: def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
    
    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)
docs: 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.