def MultiplicativeModularSparseDense(sparsity, d_feature):
  """Returns a replacement of Dense layer which uses less parameters.

  The layer uses number of modules equal to `sparsity`. It is a combination of
  multiplicative dense and locally connected dense layers.

  Args:
    sparsity: The sparsity of the layer; the output vector is divided into this
        number of modules.
    d_feature: Dimensionality of input and output tensor.
  """

  assert d_feature % sparsity == 0
  d_module = d_feature // sparsity

  return tl.Serial(
      # Weight below is used for per-head preprocessing of an embedding.
      tl.Weights(init.RandomNormalInitializer(stddev=0.5),
                 shape=[sparsity, d_feature]),
      # Weight below is a kernel of multiplicative dense, shared across heads.
      tl.Weights(init.GlorotUniformInitializer(), [d_feature, d_module]),
      # Weight below is a kernel of modular dense.
      tl.Weights(functools.partial(init.GlorotUniformInitializer(),
                                   nonreceptive_dims=[0]),
                 [sparsity, d_module, d_module]),
      # To save memory the per-head preprocessing and multiplying by
      # kernels is done in a single einsum.
      tl.Fn('SparseDenseEinsum',
            (lambda kmod, kmult, multiplier, embeds:  # pylint: disable=g-long-lambda
             jnp.einsum('hxo,dx,hd,...d->...ho', kmod, kmult, multiplier, embeds
                        ))),
      MergeLastTwoAxes(),
      # Weight below is bias after dense, per-head.
      tl.Weights(init.RandomNormalInitializer(1e-6), [d_feature]),
      tl.Add(),
      )
Returns a replacement of Dense layer which uses less parameters.

  The layer uses number of modules equal to `sparsity`. It is a combination of
  multiplicative dense and locally connected dense layers.

  Args:
    sparsity: The sparsity of the layer; the output vector is divided into this
        number of modules.
    d_feature: Dimensionality of input and output tensor.
  
trax/layers/research/sparsity.py