Documentaion suggestion

[david-waterworth]

This is a suggestion rather than a bug. Each module has an apply function which is invoke by the framework when __new__ is called. A side effect of this is intellisense doesn't work very well since the detailed comments on the apply method aren't shown, instead the single line module comment is displayed. i.e. for

x = flax.nn.normalization.BatchNorm(x)

All I see is "BatchNorm Module.". It may be more useful to move the comments from apply up to the module level. i.e.

class BatchNorm(base.Module):
    """Normalizes the input using batch statistics.

    Args:
      x: the input to be normalized.
      batch_stats: a `flax.nn.Collection` used to store an exponential moving
        average of the batch statistics (default: None).
      use_running_average: if true, the statistics stored in batch_stats
        will be used instead of computing the batch statistics on the input.
      axis: the feature or non-batch axis of the input.
      momentum: decay rate for the exponential moving average of
        the batch statistics.
      epsilon: a small float added to variance to avoid dividing by zero.
      dtype: the dtype of the computation (default: float32).
      bias:  if True, bias (beta) is added.
      scale: if True, multiply by scale (gamma).
        When the next layer is linear (also e.g. nn.relu), this can be disabled
        since the scaling will be done by the next layer.
      bias_init: initializer for bias, by default, zero.
      scale_init: initializer for scale, by default, one.
      axis_name: the axis name used to combine batch statistics from multiple
        devices. See `jax.pmap` for a description of axis names (default: None).

    Returns:
      Normalized inputs (this same shape as inputs).
    """

[avital]

Good observation. The problem with your proposal is that now a linter would correctly complain that the class shouldn't have arguments in its docstring (and where is the docstring for apply?)

To me this is another bit of evidence in favor of #208

[avital]

Consensus on the team is to land #208, which will fix this issue.