Multi-axis batch normalization support and changing from num_virtual_batches to virtual_batch_size

Key changes:
1) adding support for multi-axis batch norm by allowing the axis to be either an int or a list of ints
2) multi-axis batch norm is handled entirely by TensorFlow ops at the moment (no special kernel) and the performance is heavily dependent on the Tensor format (see reduce_* kernels for reduction rules)
3) Fix ghost batch norm by sharing the same gamma/beta/mean/var parameter across all virtual batches
4) Change ghost batch norm API to virtual_batch_size to be more consistent with its intended use case.

PiperOrigin-RevId: 172012360