Revert "[MTIA] (3/n) Implement PyTorch APIs to query/reset device pea…

…k memory usage (pytorch#143347)"

This reverts commit efe21ee.

Reverted pytorch#143347 on behalf of https://github.com/huydhn due to D67118173 has been backed out internally ([comment](pytorch#143347 (comment)))

docs/source/mtia.memory.rst

@@ -11,4 +11,3 @@ The MTIA backend is implemented out of the tree, only interfaces are be defined
     :nosignatures:
 
     memory_stats
-    max_memory_allocated

docs/source/mtia.rst

@@ -19,7 +19,6 @@ The MTIA backend is implemented out of the tree, only interfaces are be defined
     is_available
     is_initialized
     memory_stats
-    max_memory_allocated
     get_device_capability
     empty_cache
     record_memory_history

torch/mtia/memory.py

@@ -10,30 +10,29 @@
 from ._utils import _get_device_index
 
 
-def max_memory_allocated(device: Optional[_device_t] = None) -> int:
-    r"""Return the maximum memory allocated in bytes for a given device.
+def memory_stats(device: Optional[_device_t] = None) -> Dict[str, Any]:
+    r"""Return a dictionary of MTIA memory allocator statistics for a given device.
 
     Args:
         device (torch.device, str, or int, optional) selected device. Returns
             statistics for the current device, given by current_device(),
             if device is None (default).
     """
     if not is_initialized():
-        return 0
-    return memory_stats(device).get("dram", 0).get("peak_bytes", 0)
+        return {}
+    return torch._C._mtia_memoryStats(_get_device_index(device, optional=True))
 
 
-def memory_stats(device: Optional[_device_t] = None) -> Dict[str, Any]:
-    r"""Return a dictionary of MTIA memory allocator statistics for a given device.
+def max_memory_allocated(device: Optional[_device_t] = None) -> int:
+    r"""Return the maximum memory allocated in bytes for a given device.
 
     Args:
-        device (torch.device, str, or int, optional) selected device. Returns
-            statistics for the current device, given by current_device(),
-            if device is None (default).
+        device (torch.device or int, optional): selected device. Returns
+            statistic for the current device, given by :func:`~torch.mtia.current_device`,
+            if :attr:`device` is ``None`` (default).
     """
-    if not is_initialized():
-        return {}
-    return torch._C._mtia_memoryStats(_get_device_index(device, optional=True))
+
+    return memory_stats(device=device).get("allocated_bytes.all.peak", 0)
 
 
 __all__ = [