def SparseDenseWithOptions(n_units, d_input=None, sparsity_type=None,
                           sparsity=0, d_lowrank=None, prob_sparse=None,
                           mode=None, use_bias=True, use_bfloat16=False):
  """Configurable sparse version of Dense layer."""
  if prob_sparse is not None:
    if mode is not None and mode != 'train':
      # For non-training modes, we want to use a sparse variant.
      # This is different than simply prob_sparse being None, as the weights of
      # the model are different.
      prob_sparse = 1.0
    return RandomLayer(
        SparseDenseWithOptions(n_units, d_input, sparsity_type, sparsity,
                               d_lowrank, use_bias=use_bias,
                               use_bfloat16=use_bfloat16),
        tl.Dense(n_units, use_bias=use_bias, use_bfloat16=use_bfloat16),
        prob_sparse)

  if sparsity_type is None or sparsity_type == 'None' or sparsity == 0:
    return tl.Dense(n_units, use_bias=use_bias, use_bfloat16=use_bfloat16)
  if sparsity_type == 'mult':
    return FactoredDense(sparsity, d_input, n_units, use_bias=use_bias,
                         use_bfloat16=use_bfloat16)

  assert not use_bfloat16  # use_bfloat16 is unsupported for other variants
  if sparsity_type == 'lowrank':
    assert use_bias  # use_bias=False is unsupported
    return LowRankDense(n_units, d_lowrank)
  if sparsity_type == 'einsum':
    return EinsumDense(d_input, n_units, use_bias=use_bias)
  if sparsity_type == 'local':
    assert use_bias  # use_bias = False is unsupported
    assert n_units % sparsity == 0
    return LocallyConnectedDense(sparsity, n_units/sparsity)
  if sparsity_type == 'local3':
    assert use_bias  # use_bias = False is unsupported
    assert n_units % sparsity == 0
    return LocallyConnectedDense(sparsity, n_units/sparsity, kernel_size=3)

  raise ValueError('Unknown sparsity type: {}'.format(sparsity_type))
Configurable sparse version of Dense layer.
trax/layers/research/sparsity.py