bfloat16 numeric stability w/ log_softmax + cross_entropy as based on Flax examples

### Problem you have encountered:

Testing my EfficientNet impl on TPU w/ bfloat16 the training collapsed 2/3 of the way through my training schedule. The models were training reasonably on FP32 GPU, and seemed to be better behaved w/ float16 GPU and the dynamic scaler (although seems the fix here improves late results in that case as well after restoring from an earlier checkpoint).

I went through my train code (https://github.com/rwightman/efficientnet-jax/blob/master/tf_linen_train.py) based on ImageNet example here, and RMSProp impl (https://github.com/rwightman/efficientnet-jax/blob/master/jeffnet/linen/optim/rmsprop_tensorflow.py) couldn't find any issues.

I ended up fixing the issue by altering the loss fn. Instead of leaving the dtype as bfloat16 for the log_softmax I cast to float32. It is kept in float32 for the onehot and subsequent sum/mean reduction. I don't see this pattern (casting to float32) done in any Flax examples involving log_softmax + xent but it seems like it's something that will crop up for others. I did try leaving log_softmax in bfloat16 and just casting to float32 for the loss reductions. That (may have) reduced the speed of collapse but didn't prevent it.

I've tested the above fix by restoring from a point in my train schedule where the loss was continually increasing and accuracy metrics were falling off the cliff. Everything zipped back to normal. Training from scratch since has been going well.

My loss impl here (https://github.com/rwightman/efficientnet-jax/blob/master/jeffnet/common/loss.py). I keep the softmax outside of my networks.

Another somewhat related issue. I'd already cast to float32 for avgpool means as that was causing issues on GPU mixed-prec. It should be noted that all the ops mentioned here (reductions, log_softmax) are auto cast to fp32 when using AMP in PyTorch on GPU.

### What you expected to happen:

Training to converge to similar results as EfficientNet paper, Google TF TPU impl, and my own PyTorch impl.

### Steps to reproduce:

Run train script linked above. Hparams can be provided.