Pass axis=dim as a keyword arg (minor cleanup)

For consistency with the following `var` call, and because `axis` is
officially a keyword argument, change the `mean` call to pass `axis=dim`
as a keyword argument instead of a positional argument.

Issue: #7542

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -3034,7 +3034,7 @@ def _aten_var_mean_correction(tensor, dim=None, correction=1, keepdim=False):
   # be nullable, we still need to check for `None` per the API contract.
   if correction is None:
     correction = 1
-  mean = jnp.mean(tensor, dim, keepdims=keepdim)
+  mean = jnp.mean(tensor, axis=dim, keepdims=keepdim)
   # TODO: Pass in the `mean=mean` argument once `jax.numpy.var` supports it.
   var = jnp.var(tensor, axis=dim, ddof=correction, keepdims=keepdim)
   return var, mean