def LayerNormSquash(mode, width=128):  # pylint: disable=invalid-name
  """Dense-LayerNorm-Tanh normalizer inspired by ACME."""
  # https://github.com/deepmind/acme/blob/master/acme/jax/networks/continuous.py#L34
  del mode
  return tl.Serial([
      tl.Dense(width),
      tl.LayerNorm(),
      tl.Tanh(),
  ])
Dense-LayerNorm-Tanh normalizer inspired by ACME.
trax/rl/normalization.py