Add functions from `torch.masked._ops` to `__all__` for `torch.masked` (#131288)
Add the non-private operations imported in this file to `__all__` so that pyright considers them to be publicly exported. Solves this error:
```
"mean" is not exported from module "torch.masked" Pylance[reportPrivateImportUsage]
```
Related: https://github.com/pytorch/pytorch/pulls?q=pyright+export
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131288
Approved by: https://github.com/ezyang
diff --git a/test/allowlist_for_publicAPI.json b/test/allowlist_for_publicAPI.json
index b91e19a..b8822b2 100644
--- a/test/allowlist_for_publicAPI.json
+++ b/test/allowlist_for_publicAPI.json
@@ -476,6 +476,31 @@
"vecdot",
"vector_norm"
],
+ "torch.masked": [
+ "amax",
+ "amin",
+ "argmax",
+ "argmin",
+ "as_masked_tensor",
+ "cumprod",
+ "cumsum",
+ "is_masked_tensor",
+ "log_softmax",
+ "logaddexp",
+ "logsumexp",
+ "masked_tensor",
+ "MaskedTensor",
+ "mean",
+ "median",
+ "norm",
+ "normalize",
+ "prod",
+ "softmax",
+ "softmin",
+ "std",
+ "sum",
+ "var"
+ ],
"torch.multiprocessing": [
"Array",
"AuthenticationError",
diff --git a/torch/masked/__init__.py b/torch/masked/__init__.py
index 18d1b9f..d00ba1e 100644
--- a/torch/masked/__init__.py
+++ b/torch/masked/__init__.py
@@ -31,8 +31,27 @@
__all__ = [
+ "amax",
+ "amin",
+ "argmax",
+ "argmin",
"as_masked_tensor",
+ "cumprod",
+ "cumsum",
"is_masked_tensor",
+ "log_softmax",
+ "logaddexp",
+ "logsumexp",
"masked_tensor",
"MaskedTensor",
+ "mean",
+ "median",
+ "norm",
+ "normalize",
+ "prod",
+ "softmax",
+ "softmin",
+ "std",
+ "sum",
+ "var",
]
