Improve Landing Page example

[cgarciae]

@andsteing suggests:

• mention what nn is
• use shape inference
• use dataclass attributes
• personally I find "inline submodules" a bit cryptic

maybe something like this?
(feel free to change)

    import flax.linen as nn

    class MLP(nn.Module):
      out_dims: int
      @nn.compact
      def __call__(self, x):
        x = x.reshape([x.shape[0], -1])    # shape inference
        x = nn.Dense(128)(x)               # fully connected layer
        x = nn.relu(x)
        x = nn.Dense(self.out_dims)(x)     # dataclass attributes
        return x

    model = MLP(out_dims=10)               # create model

    x = jnp.ones((1, 28, 28, 1))           # fake data
    variables = model.init(PRNGKey(42), x) # initialize weights
    y = model.apply(variables, x)          # make forward pass

[Chuxiaof]

Nitpicking:

1. It's probably more conventional to use (x.shape[0], -1) instead of [x.shape[0], -1] as the argument to reshape.
2. Maybe use x = jnp.empty((1, 28, 28, 1)) to stress the argument to model.init is fake data?
3. An empty line after out_dims: int would like great.

[cgarciae]

@Chuxiaof thanks for the feedback, will include your suggestions :)
@andsteing Maybe "create submodules inline" or "compact submodules"?

[cgarciae]

How about:

   class MLP(nn.Module):                    # dataclass Modules
     out_dims: int                         

     @nn.compact
     def __call__(self, x):
       x = jnp.reshape(x, (x.shape[0], -1)) # good ol' numpy api
       x = nn.Dense(128)(x)                 # create submodules inline
       x = nn.relu(x)
       x = nn.Dense(self.out_dims)(x)       # shape inference
       return x

   model = MLP(out_dims=10)                 # create model

   x = jnp.empty((4, 28, 28, 1))            # fake data
   variables = model.init(PRNGKey(42), x)   # initialize weights
   y = model.apply(variables, x)            # make forward pass

[cgarciae]

Should we should just remove the comment over reshape about the numpy api?

[Chuxiaof]

The example looks great! I think the comment is fine, but it might be clearer to say something like "jax.numpy implements most of numpy's API".