Repository: pytorch/pytorch

Issue:
repeated warning `UserWarning: TypedStorage is deprecated`
### 🐛 Describe the bug

I noticed that the warning `UserWarning: TypedStorage is deprecated` that results from calling `Tensor.storage()` is issued repeatedly, even though the documentation about `torch.set_warn_always` says that by default pytorch warnings should only be issued once per process. While modifying my code to call `untyped_storage()` instead of `storage()` could get rid of the excessive warnings, I would prefer not to since I want to keep my code compatible with pre-2.0.0 versions for now, and `untyped_storage()` is only available >=2.0.0.

Example:
```
import torch

torch.tensor([1,2,3]).storage()
torch.tensor([1,2,3]).storage()
torch.tensor([1,2,3]).storage()
```

Output:
```
/path_to_file/test_warning.py:3: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  torch.tensor([1,2,3]).storage()
/path_to_file/test_warning.py:4: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  torch.tensor([1,2,3]).storage()
/path_to_file/test_warning.py:5: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  torch.tensor([1,2,3]).storage()
```

### Versions

```
Collecting environment information...
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.2.1 (arm64)
GCC version: Could not collect
Clang version: 12.0.0 (clang-1200.0.32.28)
CMake version: version 3.26.0
Libc version: N/A

Python version: 3.9.16 (main, Mar  1 2023, 12:19:04)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.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] numpy==1.24.2
[pip3] torch==2.0.0
[conda] numpy                     1.24.2                   pypi_0    pypi
[conda] torch                     2.0.0                    pypi_0    pypi
```

cc @albanD


Discussion:
cc @kurtamohler could you fix that?

As a temporary workaround, you can silent this warning with regular python warning filter: `python -W ignore::UserWarning: foo.py`
The complication is that the warning message contains the line of code that triggered the warning, with `warnings.warn(..., stacklevel=...)`:

https://github.com/pytorch/pytorch/blob/0b094ca37f6ac0d3c1899653ea069dbaaf33176b/torch/storage.py#L365-L372

So one option is to remove the `stacklevel` arg. But having the stack in the message is helpful for finding where the warning came from (https://github.com/pytorch/pytorch/pull/89867)

Another option is to only emit the warning the first time `_warn_typed_storage_removal()` gets called. For testing purposes, we could have a flag to make it always warn. Does that seem reasonable?

cc @ezyang 
Yeah, let's just have this warn only once and never again. If you're trying to scrub these, you can comment it out

Resolving patch:
diff --git a/torch/storage.py b/torch/storage.py
--- a/torch/storage.py
+++ b/torch/storage.py
@@ -306,14 +306,37 @@ def _isint(x):
     else:
         return isinstance(x, int)
 
+_always_warn_typed_storage_removal = False
+
+def _get_always_warn_typed_storage_removal():
+    return _always_warn_typed_storage_removal
+
+def _set_always_warn_typed_storage_removal(always_warn):
+    global _always_warn_typed_storage_removal
+    assert isinstance(always_warn, bool)
+    _always_warn_typed_storage_removal = always_warn
+
 def _warn_typed_storage_removal(stacklevel=2):
-    message = (
-        "TypedStorage is deprecated. It will be removed in the future and "
-        "UntypedStorage will be the only storage class. This should only matter "
-        "to you if you are using storages directly.  To access UntypedStorage "
-        "directly, use tensor.untyped_storage() instead of tensor.storage()"
-    )
-    warnings.warn(message, UserWarning, stacklevel=stacklevel + 1)
+    global _always_warn_typed_storage_removal
+
+    def is_first_time():
+        if not hasattr(_warn_typed_storage_removal, 'has_warned'):
+            return True
+        else:
+            return not _warn_typed_storage_removal.__dict__['has_warned']
+
+    if _get_always_warn_typed_storage_removal() or is_first_time():
+        message = (
+            "TypedStorage is deprecated. It will be removed in the future and "
+            "UntypedStorage will be the only storage class. This should only matter "
+            "to you if you are using storages directly.  To access UntypedStorage "
+            "directly, use tensor.untyped_storage() instead of tensor.storage()"
+        )
+        warnings.warn(message, UserWarning, stacklevel=stacklevel + 1)
+        _warn_typed_storage_removal.__dict__['has_warned'] = True
+
+def _reset_warn_typed_storage_removal():
+    _warn_typed_storage_removal.__dict__['has_warned'] = False
 
 class TypedStorage:
     is_sparse = False


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -37,7 +37,8 @@
     skipCUDAMemoryLeakCheckIf, BytesIOContext,
     skipIfRocm, skipIfNoSciPy, TemporaryFileName, TemporaryDirectoryName,
     wrapDeterministicFlagAPITest, DeterministicGuard, CudaSyncGuard,
-    skipIfNotRegistered, bytes_to_scalar, parametrize, skipIfMps, noncontiguous_like)
+    skipIfNotRegistered, bytes_to_scalar, parametrize, skipIfMps, noncontiguous_like,
+    AlwaysWarnTypedStorageRemoval)
 from multiprocessing.reduction import ForkingPickler
 from torch.testing._internal.common_device_type import (
     expectedFailureMeta,
@@ -6714,15 +6715,39 @@ def test_typed_storage_deprecation_warning(self):
         # Check that each of the TypedStorage function calls produce a warning
         # if warnings are reset between each
         for f in funcs:
-            with warnings.catch_warnings(record=True) as w:
-                warnings.resetwarnings()
-                f()
-                self.assertEqual(len(w), 1, msg=str([str(a) for a in w]))
-                warning = w[0].message
-                self.assertTrue(warning, DeprecationWarning)
-                self.assertTrue(re.search(
-                    '^TypedStorage is deprecated',
-                    str(warning)))
+            with AlwaysWarnTypedStorageRemoval(True):
+                with warnings.catch_warnings(record=True) as w:
+                    warnings.resetwarnings()
+                    f()
+                    self.assertEqual(len(w), 1, msg=str([str(a) for a in w]))
+                    warning = w[0].message
+                    self.assertTrue(warning, DeprecationWarning)
+                    self.assertTrue(re.search(
+                        '^TypedStorage is deprecated',
+                        str(warning)))
+
+        # Test that only the first warning is raised by default
+        torch.storage._reset_warn_typed_storage_removal()
+        with warnings.catch_warnings(record=True) as w:
+            warnings.resetwarnings()
+            torch.FloatStorage()
+            torch.randn(10).storage()
+            self.assertEqual(len(w), 1, msg=str([str(a) for a in w]))
+            warning = w[0].message
+            self.assertTrue(re.search(
+                '^TypedStorage is deprecated',
+                str(warning)))
+            # Check the line of code from the warning's stack
+            with open(w[0].filename, encoding="utf-8") as f:
+                code_line = f.readlines()[w[0].lineno - 1]
+            self.assertTrue(re.search(re.escape('torch.FloatStorage()'), code_line))
+
+        # Check that warnings are not emitted if it happened in the past
+        with warnings.catch_warnings(record=True) as w:
+            warnings.resetwarnings()
+            torch.FloatStorage()
+            torch.randn(10).storage()
+            self.assertEqual(len(w), 0, msg=str([str(a) for a in w]))
 
     def test_from_file(self):
         def assert_with_filename(filename):
diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
--- a/torch/testing/_internal/common_utils.py
+++ b/torch/testing/_internal/common_utils.py
@@ -1123,6 +1123,18 @@ def __exit__(self, exception_type, exception_value, traceback):
             self.deterministic_restore,
             warn_only=self.warn_only_restore)
 
+class AlwaysWarnTypedStorageRemoval:
+    def __init__(self, always_warn):
+        assert isinstance(always_warn, bool)
+        self.always_warn = always_warn
+
+    def __enter__(self):
+        self.always_warn_restore = torch.storage._get_always_warn_typed_storage_removal()
+        torch.storage._set_always_warn_typed_storage_removal(self.always_warn)
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        torch.storage._set_always_warn_typed_storage_removal(self.always_warn_restore)
+
 # Context manager for setting cuda sync debug mode and reset it
 # to original value
 # we are not exposing it to the core because sync debug mode is