def FactoredDense(n_modules, d_in, d_out, use_bias=True, use_bfloat16=False):
  r"""Returns a Dense-like layer, internally factored to use fewer parameters.

  This layer treats an activation vector as if divided into :math:`M`
  subvectors (``n_modules`` 'modules'). It uses this factored view to compute
  a :py:class:`Dense`-like mapping with high mixing/connectivity, but using
  approximately :math:`1/M` the number of weights of a similarly dimensioned
  :py:class:`Dense` layer.

  More specifically, each activation vector of dimensionality ``n_in`` is
  multiplied element-wise (a generalized form of gating) with ``n_modules``
  vectors also of dimensionality ``n_in``. The resulting vectors are projected
  to the subvector/module dimensionality ``d_out / n_modules`` via a matrix
  multiply, and finally reshaped back to a single vector of dimensionality
  ``d_out``. Optionally, a bias vector of dimensionality ``d_out`` is added at
  the end. All the above-mentioned non-input objects -- gating vectors,
  projection matrix, and optional bias -- are trainable weights.

  Args:
    n_modules: Number by which an activation vector is divided into subvectors
        (modules) for the factored computation.
    d_in: Last/innermost dimension of input array.
    d_out: Last/innermost dimension of output array.
    use_bias: If True, add bias vectors at the end of the layer; else end the
        layer with the matrix multiply.
    use_bfloat16: If True, use bfloat16 weights; else use float32 weights.
  """
  if d_out % n_modules != 0:
    raise ValueError(f'Value d_out ({d_out}) must be a multiple of arg '
                     f'n_modules ({n_modules}).')
  d_module = d_out // n_modules

  def GatingVectors():
    return tl.Weights(init.RandomNormalInitializer(stddev=0.5),
                      shape=[n_modules, d_in],
                      use_bfloat16=use_bfloat16)

  def ProjectionMatrix():
    return tl.Weights(init.GlorotUniformInitializer(),
                      shape=[d_in, d_module],
                      use_bfloat16=use_bfloat16),

  def Bias():
    return tl.Weights(init.RandomNormalInitializer(1e-6),
                      shape=[d_out],
                      use_bfloat16=use_bfloat16),

  layers = [
      GatingVectors(),
      ProjectionMatrix(),
      _GateAndProject(),
      MergeLastTwoAxes(),
  ]
  if use_bias:
    layers += [Bias(), tl.Add()]

  return tl.Serial(layers)
Returns a Dense-like layer, internally factored to use fewer parameters.

  This layer treats an activation vector as if divided into :math:`M`
  subvectors (``n_modules`` 'modules'). It uses this factored view to compute
  a :py:class:`Dense`-like mapping with high mixing/connectivity, but using
  approximately :math:`1/M` the number of weights of a similarly dimensioned
  :py:class:`Dense` layer.

  More specifically, each activation vector of dimensionality ``n_in`` is
  multiplied element-wise (a generalized form of gating) with ``n_modules``
  vectors also of dimensionality ``n_in``. The resulting vectors are projected
  to the subvector/module dimensionality ``d_out / n_modules`` via a matrix
  multiply, and finally reshaped back to a single vector of dimensionality
  ``d_out``. Optionally, a bias vector of dimensionality ``d_out`` is added at
  the end. All the above-mentioned non-input objects -- gating vectors,
  projection matrix, and optional bias -- are trainable weights.

  Args:
    n_modules: Number by which an activation vector is divided into subvectors
        (modules) for the factored computation.
    d_in: Last/innermost dimension of input array.
    d_out: Last/innermost dimension of output array.
    use_bias: If True, add bias vectors at the end of the layer; else end the
        layer with the matrix multiply.
    use_bfloat16: If True, use bfloat16 weights; else use float32 weights.
  
trax/layers/research/sparsity.py