def LocallyConvDense(n_modules, n_units, mode, kernel_size=1,
                     length_kernel_size=1):
  """Layer using local convolutions for approximation of Dense layer.

  The layer splits the last axis of a tensor into `n_modules`, then runs
  a convolution on all those modules, and concatenates their results.
  It is similar to LocallyConnectedDense above, but shares weights.

  Args:
    n_modules: Indicates how many modules (pixels) should be input and output
        split into for processing.
    n_units: how many outputs (filters) should each module generate.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
    kernel_size: The size of the kernel to be used.
    length_kernel_size: If > 1, also do causal convolution on the previous axis,
      which is often the sentence length in sequence models.

  Returns:
      LocallyConvDense base.Layer.
  """
  if n_modules == 1:
    return tl.Dense(n_units)
  if kernel_size % 2 != 1:
    raise ValueError('Currently we only handle odd kernel sizes.')
  half = (kernel_size - 1) // 2
  pad_widths = [[0, 0], [0, 0], [half, half], [0, 0]]
  return tl.Serial(
      tl.SplitLastAxis(n_modules),
      tl.Fn('Pad', lambda x: jnp.pad(x, pad_width=pad_widths, mode='constant')),
      _RememberPad(length_kernel_size-1, mode=mode),
      tl.Conv(n_units, kernel_size=(length_kernel_size, kernel_size)),
      tl.MergeLastTwoAxes()
  )
Layer using local convolutions for approximation of Dense layer.

  The layer splits the last axis of a tensor into `n_modules`, then runs
  a convolution on all those modules, and concatenates their results.
  It is similar to LocallyConnectedDense above, but shares weights.

  Args:
    n_modules: Indicates how many modules (pixels) should be input and output
        split into for processing.
    n_units: how many outputs (filters) should each module generate.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
    kernel_size: The size of the kernel to be used.
    length_kernel_size: If > 1, also do causal convolution on the previous axis,
      which is often the sentence length in sequence models.

  Returns:
      LocallyConvDense base.Layer.
  
trax/layers/research/sparsity.py