NNX `_compute_stats` function missing `use_fast_variance` and `mask` argument

[Marcus Chiam]

NNX _compute_stats function is missing use_fast_variance and mask argument, compared to the Linen equivalent. Is this intentional or should we add this into NNX?

cc: @cgarciae

[Cristian Garcia]

This is not intentional, most likely it was added after the initial port of the code.
If you can add use_fast_variance it would amazing!