code: 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),
      )
  
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: NeuralGPU(d_feature, steps, vocab_size, mode) and returns: tl