Repository: pytorch
Issue: Error computing the norm of MaskedTensor
Labels: triaged,module: masked operators

### 🐛 Describe the bug

According to the [documentation](https://pytorch.org/docs/2.1/masked.html#reductions), the `norm` reduction is supported for the beta `MaskedTensor`. However the following does not work:

```python
>>> import torch
>>> x = torch.rand(2, 3, 4, 5)
>>> torch.masked.as_masked_tensor(x, x>0.5).norm()
/opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/masked/maskedtensor/core.py:156: UserWarning: The PyTorch API of MaskedTensors is in prototype stage and will change in the near future. Please open a Github issue for features requests and see our documentation on the torch.masked module for further information about the project.
  warnings.warn(("The PyTorch API of MaskedTensors is in prototype stage "
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[3], line 1
----> 1 torch.masked.as_masked_tensor(x, x>0.5).norm()

File /opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/_tensor.py:705, in Tensor.norm(self, p, dim, keepdim, dtype)
    703 r"""See :func:`torch.norm`"""
    704 if has_torch_function_unary(self):
--> 705     return handle_torch_function(
    706         Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype
    707     )
    708 return torch.norm(self, p, dim, keepdim, dtype=dtype)

File /opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/overrides.py:1577, in handle_torch_function(public_api, relevant_args, *args, **kwargs)
   1571     warnings.warn("Defining your `__torch_function__ as a plain method is deprecated and "
   1572                   "will be an error in future, please define it as a classmethod.",
   1573                   DeprecationWarning)
   1575 # Use `public_api` instead of `implementation` so __torch_function__
   1576 # implementations can do equality/identity comparisons.
-> 1577 result = torch_func_method(public_api, types, args, kwargs)
   1579 if result is not NotImplemented:
   1580     return result

File /opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/masked/maskedtensor/core.py:269, in MaskedTensor.__torch_function__(cls, func, types, args, kwargs)
    267 from ._ops_refs import _MASKEDTENSOR_FUNCTION_TABLE
    268 if func in _MASKEDTENSOR_FUNCTION_TABLE:
--> 269     return _MASKEDTENSOR_FUNCTION_TABLE[func](*args, **kwargs)
    271 if not all(issubclass(cls, t) for t in types):
    272     return NotImplemented

File /opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/masked/maskedtensor/_ops_refs.py:184, in _general_function_reductions(func, *args, **kwargs)
    182 @register_function_func(NATIVE_REDUCE_FNS + TORCH_REDUCE_FNS + TENSOR_REDUCE_FNS)
    183 def _general_function_reductions(func, *args, **kwargs):
--> 184     return _apply_reduction(func, *args, **kwargs)

File /opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/masked/maskedtensor/reductions.py:172, in _apply_reduction(fn, *args, **kwargs)
    170     return TORCH_REDUCE_MAP[fn](*args, **kwargs)
    171 if fn in TENSOR_REDUCE_MAP:
--> 172     return TENSOR_REDUCE_MAP[fn](*args, **kwargs)
    173 return NotImplemented

File /opt/homebrew/Caskroom/miniforge/base/envs/torch/lib/python3.10/site-packages/torch/masked/maskedtensor/reductions.py:128, in _torch_grad_reduce.<locals>.grad_reduce(*args, **kwargs)
    126     return _torch_reduce_all(fn)(args[0])
    127 # TODO: autograd.Function doesn't support kwarg
--> 128 input, dim, keepdim, dtype = _reduce_dim_args(*args, **kwargs)
    129 return _torch_reduce_dim(fn)(input, dim, keepdim, dtype)

TypeError: _reduce_dim_args() got an unexpected keyword argument 'p'
```

### Versions

Tested on CPU with macOS arm64 and linux x64:

<details>
  <summary>macOS:</summary>

  ```
PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.2.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.27.7
Libc version: N/A

Python version: 3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:41:22) [Clang 13.0.1 ] (64-bit runtime)
Python platform: macOS-14.2.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] numpy==1.23.3
[pip3] onnx==1.14.0
[pip3] onnxruntime==1.14.1
[pip3] pytorch-lightning==2.0.0
[pip3] torch==2.1.2
[pip3] torch-receptive-field==0.0.1
[pip3] torchaudio==0.12.1
[pip3] torchmetrics==0.11.4
[pip3] torchvision==0.16.2
[conda] numpy                     1.23.3          py310h5d7c261_0    conda-forge
[conda] pytorch-lightning         2.0.0                    pypi_0    pypi
[conda] torch                     2.1.2                    pypi_0    pypi
[conda] torch-receptive-field     0.0.1                    pypi_0    pypi
[conda] torchaudio                0.12.1                py310_cpu    pytorch
[conda] torchmetrics              0.11.4                   pypi_0    pypi
[conda] torchvision               0.16.2                   pypi_0    pypi
  ```
</details>



<details>
  <summary>Linux:</summary>

  ```
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Rocky Linux release 8.8 (Green Obsidian) (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-18)
Clang version: 15.0.7 (Red Hat 15.0.7-1.module+el8.8.0+1144+0a4e73bd)
CMake version: version 3.26.3
Libc version: glibc-2.28

Python version: 3.10.11 (main, Apr 20 2023, 19:02:41) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.18.0-477.27.1.el8_8.x86_64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A40
Nvidia driver version: 545.23.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              16
On-line CPU(s) list: 0-15
Thread(s) per core:  1
Core(s) per socket:  16
Socket(s):           1
NUMA node(s):        1
Vendor ID:           AuthenticAMD
CPU family:          23
Model:               49
Model name:          AMD EPYC 7302P 16-Core Processor
Stepping:            0
CPU MHz:             3000.000
CPU max MHz:         3000.0000
CPU min MHz:         1500.0000
BogoMIPS:            6000.04
Virtualization:      AMD-V
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            16384K
NUMA node0 CPU(s):   0-15
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] onnx==1.14.0
[pip3] onnxruntime==1.15.0
[pip3] onnxruntime-gpu==1.15.0
[pip3] pytorch-lightning==2.0.2
[pip3] torch==2.1.2
[pip3] torchmetrics==1.0.1
[pip3] torchview==0.2.6
[pip3] torchvision==0.16.2
[pip3] triton==2.1.0
[conda] numpy                     1.24.3                   pypi_0    pypi
[conda] pytorch-lightning         2.0.2                    pypi_0    pypi
[conda] torch                     2.1.2                    pypi_0    pypi
[conda] torchmetrics              1.0.1                    pypi_0    pypi
[conda] torchview                 0.2.6                    pypi_0    pypi
[conda] torchvision               0.16.2                   pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi
  ```

</details>