title: trax tl.Relu and tl.ShiftRight layers are nested inside Serial Combinator
body: <p>I am trying to build an attention model but Relu and ShiftRight layer by default nested inside the Serial Combinator.
  This further gives me errors in training.</p>
  <pre><code>layer_block = tl.Serial(
      tl.Relu(),
      tl.LayerNorm(), )
  
  x = np.array([[-2, -1, 0, 1, 2],
                [-20, -10, 0, 10, 20]]).astype(np.float32) 
  
  layer_block.init(shapes.signature(x)) y = layer_block(x)
  
  print(f'layer_block: {layer_block}')
  </code></pre>
  <h3>Output</h3>
  <pre><code>layer_block: Serial[
    Serial[
      Relu
    ]
    LayerNorm
  ]
  </code></pre>
  <h3>Expected Output</h3>
  <pre><code>layer_block: Serial[
    Relu
    LayerNorm
  ]
  </code></pre>
  <p>The same problem arises with tl.ShiftRight()</p>
  <p>The code above is taken from official documentation <a href="https://trax-ml.readthedocs.io/en/latest/notebooks/layers_intro.html#Layers-combine-into-layers." rel="nofollow noreferrer">Example 5</a></p>
  <p>Thanks in advance</p>