def SparseFF(
    d_ff, n_elements_in_block=32, d_lowrank=64, temperature=0.1, quant_prob=0.3,
    use_bfloat16=False, big_weights_in_bfloat16=False, mode='train',
    kernel_initializer=init.GlorotUniformInitializer(),
    bias_initializer=init.RandomNormalInitializer(1e-6),
    dropout_rate=0.0, dropout_shared_axes=None, ff_chunk_size=0,
    multiply_by_controller_output=False, kernel_scaling=False):
  """Returns Feed-forward block with sparsity.

  The original (non-sparse) FF block is a triple Dense(d_ff)-Relu-Dense
  that takes an input, makes it of size d_ff (usually larger than it was) and
  then brings it back to the original size after Relu. It is commonly used in
  Transformer models where it often accounts for most of the trainable weights.

  The original block can be slow in decoding due to the need to fetch a lot of
  weights from memory. This sparse block only allows one non-zero element
  in a block of a specified size. This is trained with straight-through Gumbel
  softmax trick.

  Args:
    d_ff: Depth/dimensionality of FeedForward layer.
    n_elements_in_block: The sparsity level. The layer is divided into blocks of
      this size, and each block has only a single element active.
    d_lowrank: The dimensionality of low-rank controller.
    temperature: The temperature of the controller during training.
    quant_prob: During training this proportion of blocks will have quantized
      mask (i.e. a single element active). The rest will use a soft mask.
    use_bfloat16: Whether to use bfloat16 for weights.
    big_weights_in_bfloat16: : Whether to use bfloat16 for main weights of the
      FeedForward layer.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
    kernel_initializer: Function that creates a matrix of (random) initial
        connection weights `W` for the layer.
    bias_initializer: Function that creates a vector of (random) initial
        bias weights `b` for the layer.
    dropout_rate: Probability for dropping an activation value.
    dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
      along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
      way to save memory and apply consistent masks to activation vectors at
      different sequence positions.
    ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks.
    multiply_by_controller_output: whether to multiply the middle activation
      layer of FF by controller output (i.e. softmax).
    kernel_scaling: Whether to scale the kernel matrix (during init) to keep the
      variance of the layer output regardless of n_elements_in_block.
  """

  if mode == 'train' or multiply_by_controller_output:
    also_return_nondiscrete_output = True
  else:
    also_return_nondiscrete_output = False
  controller = _SparseFFController(
      d_ff=d_ff, n_elements_in_block=n_elements_in_block,
      d_lowrank=d_lowrank, temperature=temperature,
      use_bfloat16=use_bfloat16, mode=mode,
      kernel_initializer=kernel_initializer,
      bias_initializer=bias_initializer,
      also_return_nondiscrete_output=also_return_nondiscrete_output)

  main = [
      _SparseFFMain(
          d_ff=d_ff, n_elements_in_block=n_elements_in_block,
          d_lowrank=d_lowrank, quant_prob=quant_prob, use_bfloat16=use_bfloat16,
          big_weights_in_bfloat16=big_weights_in_bfloat16, mode=mode,
          kernel_initializer=kernel_initializer,
          bias_initializer=bias_initializer,
          multiply_by_controller_output=multiply_by_controller_output,
          kernel_scaling=kernel_scaling),
      # quant_mask, emb
      tl.Select([1, 0]),
      # emb, quant_mask
      tl.Dropout(rate=dropout_rate, shared_axes=dropout_shared_axes, mode=mode),
      tl.Select([1, 0]),
      # quant_mask, emb
  ]

  # We will "remember" quant_mask _after_ chunking, and "recall" this same
  # quant_mask during reverse_and_grad _before_ chunking.
  remembering = _RememberInReverse(output=False)
  recalling = _RecallQuantMaskInReverse(
      remember_layer=remembering, elements=d_ff//n_elements_in_block)

  return tl.BatchLeadingAxes(tl.Serial(
      recalling,  # emb, quant_mask
      tl.Chunk(chunk_size=ff_chunk_size, layer=tl.Serial(
          # emb, quant_mask
          tl.Select((0, 1, 0)),  # emb, quant_mask, emb
          controller,  # quant_mask, mask, emb
          main,  # quant_mask, emb/output
          )),
      remembering,  # emb/output
      ))
Returns Feed-forward block with sparsity.

  The original (non-sparse) FF block is a triple Dense(d_ff)-Relu-Dense
  that takes an input, makes it of size d_ff (usually larger than it was) and
  then brings it back to the original size after Relu. It is commonly used in
  Transformer models where it often accounts for most of the trainable weights.

  The original block can be slow in decoding due to the need to fetch a lot of
  weights from memory. This sparse block only allows one non-zero element
  in a block of a specified size. This is trained with straight-through Gumbel
  softmax trick.

  Args:
    d_ff: Depth/dimensionality of FeedForward layer.
    n_elements_in_block: The sparsity level. The layer is divided into blocks of
      this size, and each block has only a single element active.
    d_lowrank: The dimensionality of low-rank controller.
    temperature: The temperature of the controller during training.
    quant_prob: During training this proportion of blocks will have quantized
      mask (i.e. a single element active). The rest will use a soft mask.
    use_bfloat16: Whether to use bfloat16 for weights.
    big_weights_in_bfloat16: : Whether to use bfloat16 for main weights of the
      FeedForward layer.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
    kernel_initializer: Function that creates a matrix of (random) initial
        connection weights `W` for the layer.
    bias_initializer: Function that creates a vector of (random) initial
        bias weights `b` for the layer.
    dropout_rate: Probability for dropping an activation value.
    dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
      along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
      way to save memory and apply consistent masks to activation vectors at
      different sequence positions.
    ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks.
    multiply_by_controller_output: whether to multiply the middle activation
      layer of FF by controller output (i.e. softmax).
    kernel_scaling: Whether to scale the kernel matrix (during init) to keep the
      variance of the layer output regardless of n_elements_in_block.
  
trax/layers/research/sparsity.py