def squeeze_layer(axis=1):
  """Layer for squeezing dimension along the axis."""
  def init_fun(rng, input_shape):
    del rng
    if axis < 0:
      raise ValueError("squeeze_layer: negative axis is not supported")
    return (input_shape[:axis] + input_shape[(axis + 1):]), ()

  def apply_fun(params, inputs, **kwargs):
    del params, kwargs
    return inputs.squeeze(axis)

  return init_fun, apply_fun