Reduce amount of modules, clean-up

[lericson]

Use existing primitives such as nn.GELU
Remove no-op modules
Update super()
Removed trailing whitespaces in code

[lericson]

I should note that this fixes a minor discrepancy in the code compared to the JAX reference, their code says (with some trivial rewriting):

x = self.encoder(x)

x = x[:, 0]

if repr_dim is not None:
  x = nn.Dense(repr_dim)(x)
  x = nn.tanh(x)

... whereas this is what this repository says:

        x = self.encoder(x)
        x = self.pre_logits(x)

        # only support cls token now
        x = x[:, 0]

        return self.head(x)

It doesn't actually matter, but the placement of the x = x[:, 0] is earlier in the reference code. This PR does the same. It just saves some cycles I guess.