def EinsumDense(d_input, d_output, use_bias):
  """Returns a reimplementation of Dense layer, using einsum.

  While this is an equivalent of a Dense layer, it seems to be faster when used
  in decoding if used with bias (see decoding_timing_test.py ).
  This layer can be removed when we understand better the reason for the
  difference in decoding speed.

  Args:
    d_input: Dimensionality of the input tensor.
    d_output: Dimensionality of the output tensor.
    use_bias: Whether to use bias.
  """
  layers = [
      tl.Weights(init.GlorotUniformInitializer(), [d_output, d_input]),
      tl.Fn('EinsumDense',
            (lambda kernel, embeds:  # pylint: disable=g-long-lambda
             jnp.einsum('xd,...d->...x', kernel, embeds)))
  ]
  if use_bias:
    layers.extend([
        tl.Weights(init.RandomNormalInitializer(1e-6), [d_output]),
        tl.Add()
    ])
  return tl.Serial(layers)
Returns a reimplementation of Dense layer, using einsum.

  While this is an equivalent of a Dense layer, it seems to be faster when used
  in decoding if used with bias (see decoding_timing_test.py ).
  This layer can be removed when we understand better the reason for the
  difference in decoding speed.

  Args:
    d_input: Dimensionality of the input tensor.
    d_output: Dimensionality of the output tensor.
    use_bias: Whether to use bias.
  
trax/layers/research/sparsity.py