(n_modules, n_units, mode, kernel_size=1,
                     length_kernel_size=1)
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.
Layer using local convolutions for approximation of Dense layer.
Layer using local convolutions for approximation of Dense layer .
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()
  )
def LocallyConvDense ( n_modules , n_units , mode , kernel_size = 1 , length_kernel_size = 1 ) : 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 ( ) )