(sparsity, d_feature)
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.
Returns a replacement of Dense layer which uses less parameters.
Returns a replacement of Dense layer which uses less parameters .
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(),
      )
def MultiplicativeModularSparseDense ( sparsity , d_feature ) : 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 ( ) , )