class_docstr: Accelerator for XLA devices, normally TPUs.
  
  .. warning::  Use of this accelerator beyond import and instantiation is experimental.
class_docstr_tok: ["Accelerator","for","XLA","devices",",","normally","TPUs",".","..","warning",":",":","Use","of","this","accelerator","beyond","import","and","instantiation","is","experimental","."]
docstr: Gets stats for the given XLA device.
  
  Args:
      device: XLA device for which to get stats
  
  Returns:
      A dictionary mapping the metrics (free memory and peak memory) to their values.
docstr_tok: ["Gets","stats","for","the","given","XLA","device",".","Args",":","device",":","XLA","device","for","which","to","get","stats","Returns",":","A","dictionary","mapping","the","metrics","(","free","memory","and","peak","memory",")","to","their","values","."]
code: def get_device_stats(self, device: _DEVICE) -> Dict[str, Any]:
      """Gets stats for the given XLA device.
  
      Args:
          device: XLA device for which to get stats
  
      Returns:
          A dictionary mapping the metrics (free memory and peak memory) to their values.
  
      """
      import torch_xla.core.xla_model as xm
  
      memory_info = xm.get_memory_info(device)
      free_memory = memory_info["kb_free"]
      peak_memory = memory_info["kb_total"] - free_memory
      return {
          "avg. free memory (MB)": free_memory,
          "avg. peak memory (MB)": peak_memory,
      }
  
file_path: pytorch-lightning/src/lightning/pytorch/accelerators/xla.py
imports: from typing import Any, Dict
  from typing_extensions import override
  from lightning.fabric.accelerators import _AcceleratorRegistry
  from lightning.fabric.accelerators.xla import XLAAccelerator
  from lightning.fabric.utilities.types import _DEVICE
  from lightning.pytorch.accelerators.accelerator import Accelerator
label_desc: Use image node_id 1 for calling the XLAAccelerator obj's underlying member method code with example usage: obj.get_device_stats(device) and returns: dict