PyTorch/XLA API

torch_xla

torch_xla.device(index: Optional[int] = None) → device

返回一个给定的 XLA 设备实例。

如果 SPMD 已启用，则返回一个包装此进程可用的所有设备的虚拟设备。

参数

index – 要返回的 XLA 设备的索引。对应于 torch_xla.devices() 中的索引。

返回

一个 XLA torch.device。

torch_xla.devices() → List[device]

返回当前进程中所有可用的设备。

返回

一个 XLA torch.devices 列表。

torch_xla.device_count() → int

返回当前进程中可寻址设备的数量。

torch_xla.sync(wait: bool = False)

启动所有待定的图操作。

参数

wait (bool) – 是否阻塞当前进程直到执行完成。

torch_xla.compile(f: Optional[Callable] = None, full_graph: Optional[bool] = False, name: Optional[str] = None, num_different_graphs_allowed: Optional[int] = None)

使用 torch_xla 的 LazyTensor 跟踪模式优化给定的模型/函数。PyTorch/XLA 将使用给定的输入跟踪函数，然后生成图来表示此函数内发生的 pytorch 操作。此图将由 XLA 编译并在加速器上执行（由张量的设备决定）。Eager 模式将被禁用于函数的已编译区域。

参数

• model (Callable) – 要优化的模块/函数，如果不传，此函数将作为上下文管理器。

• full_graph (Optional[bool]) – 此编译是否应生成单个图。如果设置为 True，并且将生成多个图，torch_xla 将抛出带有调试信息的错误并退出。

• name (Optional[name]) – 已编译程序的名称。如果未指定，则使用函数 f 的名称。此名称将用于 PT_XLA_DEBUG 消息以及 HLO/IR 转储文件。

• num_different_graphs_allowed (Optional[python:int]) – 允许的给定模型/函数的不同跟踪图的数量。如果超过此限制，将引发错误。

示例

# usage 1
@torch_xla.compile()
def foo(x):
  return torch.sin(x) + torch.cos(x)

def foo2(x):
  return torch.sin(x) + torch.cos(x)
# usage 2
compiled_foo2 = torch_xla.compile(foo2)

# usage 3
with torch_xla.compile():
  res = foo2(x)

torch_xla.manual_seed(seed, device=None)

为当前 XLA 设备设置随机数生成器的种子。

参数

• seed (python:integer) – 要设置的状态。

• device (torch.device, optional) – 需要设置 RNG 状态的设备。如果缺失，将设置默认设备种子。

runtime

torch_xla.runtime.device_type() → Optional[str]

返回当前的 PjRt 设备类型。

如果尚未配置默认设备，则选择一个默认设备

返回

设备的字符串表示。

torch_xla.runtime.local_process_count() → int

返回此主机上运行的进程数。

torch_xla.runtime.local_device_count() → int

返回此主机上的总设备数。

假定每个进程拥有相同数量的可寻址设备。

torch_xla.runtime.addressable_device_count() → int

返回此进程可见的设备数。

torch_xla.runtime.global_device_count() → int

返回所有进程/主机上的总设备数。

torch_xla.runtime.global_runtime_device_count() → int

返回所有进程/主机上的总运行时设备数，这对 SPMD 尤其有用。

torch_xla.runtime.world_size() → int

返回作业中参与的总进程数。

torch_xla.runtime.global_ordinal() → int

返回此线程在所有进程中的全局序号。

全局序号的范围是 [0, global_device_count)。全局序号不能保证与 TPU 工作进程 ID 有任何可预测的关系，也不能保证在每个主机上是连续的。

torch_xla.runtime.local_ordinal() → int

返回此线程在此主机中的本地序号。

本地序号的范围是 [0, local_device_count)。

torch_xla.runtime.get_master_ip() → str

检索运行时的 master worker IP。此函数调用特定于后端的发现 API。

返回

master worker 的 IP 地址（字符串形式）。

torch_xla.runtime.use_spmd(auto: Optional[bool] = False)

用于启用 SPMD 模式的 API。这是启用 SPMD 的推荐方式。

如果某些张量已在非 SPMD 设备上初始化，则强制进入 SPMD 模式。这意味着这些张量将在设备之间进行复制。

参数

auto (bool) – 是否启用自动分片。有关更多详细信息，请参阅 https://github.com/pytorch/xla/blob/master/docs/spmd_advanced.md#auto-sharding

torch_xla.runtime.is_spmd()

返回 SPMD 是否已设置为执行。

torch_xla.runtime.initialize_cache(path: str, readonly: bool = False)

初始化持久化编译缓存。此 API 必须在执行任何计算之前调用。

参数

• path (str) – 用于存储持久化缓存的路径。

• readonly (bool) – 此工作进程是否具有对缓存的写入访问权限。

xla_model

torch_xla.core.xla_model.xla_device(n: Optional[int] = None, devkind: Optional[str] = None) → device

返回一个给定的 XLA 设备实例。

参数

• n (python:int, optional) – 要返回的特定实例（序号）。如果指定，将返回特定的 XLA 设备实例。否则，将返回 devkind 的第一个设备。

• devkind (string..., optional) – 如果指定，则为设备类型，例如 TPU、CUDA、CPU 或自定义 PJRT 设备。已弃用。

返回

一个具有请求实例的 torch.device。

torch_xla.core.xla_model.xla_device_hw(device: Union[str, device]) → str

返回给定设备的硬件类型。

参数

device (string or torch.device) – 将映射到真实设备的 xla 设备。

返回

给定设备的硬件类型的字符串表示。

torch_xla.core.xla_model.is_master_ordinal(local: bool = True) → bool

检查当前进程是否为主序号（0）。

参数

local (bool) – 检查的是本地主序号还是全局主序号。对于多主机复制，只有一个全局主序号（主机 0，设备 0），而有 NUM_HOSTS 个本地主序号。默认值：True

返回

一个布尔值，指示当前进程是否为主序号。

torch_xla.core.xla_model.all_reduce(reduce_type: str, inputs: Union[Tensor, List[Tensor]], scale: float = 1.0, groups: Optional[List[List[int]]] = None, pin_layout: bool = True) → Union[Tensor, List[Tensor]]

对输入张量执行就地（inplace）归约操作。

参数

• reduce_type (string) – xm.REDUCE_SUM、xm.REDUCE_MUL、xm.REDUCE_AND、xm.REDUCE_OR、xm.REDUCE_MIN 和 xm.REDUCE_MAX 之一。

• inputs – 要执行 all reduce 操作的单个 torch.Tensor 或 torch.Tensor 列表。

• scale (python:float) – 归约后应用的默认缩放值。默认值：1.0

• groups (list, optional) –

  一个列表的列表，表示 all_reduce() 操作的副本组。示例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定义了两组，一组包含 [0, 1, 2, 3] 副本，另一组包含 [4, 5, 6, 7] 副本。如果为 None，则只有一组包含所有副本。

• pin_layout (bool, optional) – 是否为此通信操作固定布局。布局固定可以防止潜在的数据损坏，当参与通信的每个进程具有略有不同的程序时，但它可能会导致某些 xla 编译失败。当您看到类似“HloModule 具有布局约束的混合”的错误消息时，请取消固定布局。

返回

如果传入单个 torch.Tensor，则返回值是一个包含归约值（跨副本）的 torch.Tensor。如果传入列表/元组，此函数将对输入张量执行就地 all-reduce 操作，并返回列表/元组本身。

torch_xla.core.xla_model.all_gather(value: Tensor, dim: int = 0, groups: Optional[List[List[int]]] = None, output: Optional[Tensor] = None, pin_layout: bool = True) → Tensor

沿给定维度执行 all-gather 操作。

参数

• value (torch.Tensor) – 输入张量。

• dim (python:int) – 收集维度。默认值：0

• groups (list, optional) –

  一个列表的列表，表示 all_gather() 操作的副本组。示例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定义了两组，一组包含 [0, 1, 2, 3] 副本，另一组包含 [4, 5, 6, 7] 副本。如果为 None，则只有一组包含所有副本。

• output (torch.Tensor) – 可选输出张量。

• pin_layout (bool, optional) – 是否为此通信操作固定布局。布局固定可以防止潜在的数据损坏，当参与通信的每个进程具有略有不同的程序时，但它可能会导致某些 xla 编译失败。当您看到类似“HloModule 具有布局约束的混合”的错误消息时，请取消固定布局。

返回

一个张量，在 dim 维度上包含来自参与副本的所有值。

torch_xla.core.xla_model.all_to_all(value: Tensor, split_dimension: int, concat_dimension: int, split_count: int, groups: Optional[List[List[int]]] = None, pin_layout: bool = True) → Tensor

对输入张量执行 XLA AllToAll() 操作。

参见：https://tensorflowcn.cn/xla/operation_semantics#alltoall

参数

• value (torch.Tensor) – 输入张量。

• split_dimension (python:int) – 应该发生分割的维度。

• concat_dimension (python:int) – 应该发生连接的维度。

• split_count (python:int) – 分割计数。

• groups (list, optional) –

  一个列表的列表，表示 all_reduce() 操作的副本组。示例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定义了两组，一组包含 [0, 1, 2, 3] 副本，另一组包含 [4, 5, 6, 7] 副本。如果为 None，则只有一组包含所有副本。

• pin_layout (bool, optional) – 是否为此通信操作固定布局。布局固定可以防止潜在的数据损坏，当参与通信的每个进程具有略有不同的程序时，但它可能会导致某些 xla 编译失败。当您看到类似“HloModule 具有布局约束的混合”的错误消息时，请取消固定布局。

返回

all_to_all() 操作的结果 torch.Tensor。

torch_xla.core.xla_model.add_step_closure(closure: Callable[[...], Any], args: Tuple[Any, ...] = (), run_async: bool = False)

将一个闭包添加到将在步骤结束时运行的列表。

在模型训练的许多时候，需要打印/报告（打印到控制台、发布到 tensorboard 等）信息，这需要检查中间张量的内容。在模型代码的不同点检查不同的张量内容需要多次执行，并且通常会导致性能问题。添加一个步骤闭包将确保它将在屏障后运行，届时所有活动的张量都已物化到设备数据。活动张量将包括由闭包参数捕获的张量。因此，使用 add_step_closure() 将确保执行一次，即使队列中有多个闭包，需要检查多个张量。步骤闭包将按其入队顺序依次运行。请注意，即使使用此 API，执行也会得到优化，但建议每 N 步限制一次打印/报告事件。

参数

• closure (callable) – 要调用的函数。

• args (tuple) – 要传递给闭包的参数。

• run_async – 如果为 True，则异步运行闭包。

torch_xla.core.xla_model.wait_device_ops(devices: List[str] = [])

等待给定设备上的所有异步操作完成。

参数

devices (string..., optional) – 需要等待其异步操作的设备。如果为空，将等待所有本地设备。

torch_xla.core.xla_model.optimizer_step(optimizer: Optimizer, barrier: bool = False, optimizer_args: Dict = {}, groups: Optional[List[List[int]]] = None, pin_layout: bool = True)

执行提供的优化器步骤并跨所有设备同步梯度。

参数

• optimizer (torch.Optimizer) – 需要调用其 step() 函数的 torch.Optimizer 实例。 step() 函数将使用 optimizer_args 命名参数进行调用。

• barrier (bool, optional) – 是否在此 API 中发出 XLA 张量屏障。如果使用 PyTorch XLA 的 ParallelLoader 或 DataParallel 支持，则不需要此项，因为屏障将由 XLA 数据加载器迭代器 next() 调用发出。默认值：False

• optimizer_args (dict, optional) – optimizer.step() 调用的命名参数字典。

• groups (list, optional) –

  一个列表的列表，表示 all_reduce() 操作的副本组。示例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定义了两组，一组包含 [0, 1, 2, 3] 副本，另一组包含 [4, 5, 6, 7] 副本。如果为 None，则只有一组包含所有副本。

• pin_layout (bool, optional) – 归约梯度时是否固定布局。有关详细信息，请参阅 xm.all_reduce。

返回

由 optimizer.step() 调用返回的相同值。

示例

>>> import torch_xla.core.xla_model as xm
>>> xm.optimizer_step(self.optimizer)

torch_xla.core.xla_model.save(data: Any, file_or_path: Union[str, TextIO], master_only: bool = True, global_master: bool = False)

将输入数据保存到文件。

保存的数据将在保存前传输到 PyTorch CPU 设备，因此后续的 torch.load() 将加载 CPU 数据。在处理视图时必须小心。与其保存视图，不如建议在张量已加载并移动到其目标设备后重新创建它们。

参数

• data – 要保存的输入数据。任何嵌套的 Python 对象组合（列表、元组、集合、字典……）。

• file_or_path – 数据保存操作的目标。可以是文件路径或 Python 文件对象。如果 master_only 为 False，则路径或文件对象必须指向不同的目标，否则来自同一主机的写入将覆盖彼此。

• master_only (bool, optional) – 是否只有主设备保存数据。如果为 False，则 file_or_path 参数对于参与复制的每个序号都应为不同的文件或路径，否则同一主机上的所有副本都将写入同一位置。默认值：True

• global_master (bool, optional) – 当 master_only 为 True 时，此标志控制是每个主机的（如果 global_master 为 False）保存内容，还是只有全局主（序号 0）保存。默认值：False

示例

>>> import torch_xla.core.xla_model as xm
>>> xm.wait_device_ops() # wait for all pending operations to finish.
>>> xm.save(obj_to_save, path_to_save)
>>> xm.rendezvous('torch_xla.core.xla_model.save') # multi process context only

torch_xla.core.xla_model.rendezvous(tag: str, payload: bytes = b'', replicas: List[int] = []) → List[bytes]

等待所有 mesh 客户端到达命名集合点。

注意：PJRT 不支持 XRT mesh 服务器，因此这实际上是 xla_rendezvous 的别名。

参数

• tag (string) – 要加入的集合点的名称。

• payload (bytes, optional) – 要发送到集合点的数据。

• replicas (list, python:int) – 参与集合点的副本序号。空表示 mesh 中的所有副本。默认值：[]

返回

所有其他核心交换的数据负载，其中核心序号 i 的数据负载位于返回元组的位置 i。

示例

>>> import torch_xla.core.xla_model as xm
>>> xm.rendezvous('example')

torch_xla.core.xla_model.mesh_reduce(tag: str, data, reduce_fn: Callable[[...], Any]) → Union[Any, ToXlaTensorArena]

执行图外客户端 mesh 归约。

参数

• tag (string) – 要加入的集合点的名称。

• data – 要归约的数据。 reduce_fn 可调用函数将接收一个列表，其中包含来自所有 mesh 客户端进程（每个核心一个）的相同数据的副本。

• reduce_fn (callable) – 接收 data 类型的对象列表并返回归约结果的函数。

返回

归约后的值。

示例

>>> import torch_xla.core.xla_model as xm
>>> import numpy as np
>>> accuracy = xm.mesh_reduce('test_accuracy', accuracy, np.mean)

torch_xla.core.xla_model.set_rng_state(seed: int, device: Optional[str] = None)

设置随机数生成器状态。

参数

• seed (python:integer) – 要设置的状态。

• device (string, optional) – 需要设置 RNG 状态的设备。如果缺失，将设置默认设备种子。

torch_xla.core.xla_model.get_rng_state(device: Optional[str] = None) → int

获取当前的随机数生成器状态。

参数

device (string, optional) – 需要检索其 RNG 状态的设备。如果缺失，将设置默认设备种子。

返回

RNG 状态（整数形式）。

torch_xla.core.xla_model.get_memory_info(device: Optional[device] = None) → MemoryInfo

检索设备内存使用情况。

参数

• device – Optional[torch.device] 请求内存信息的设备。

• device. (如果不传入，将使用默认设备) –

返回

包含给定设备内存使用情况的 MemoryInfo 字典。

示例

>>> xm.get_memory_info()
{'bytes_used': 290816, 'bytes_limit': 34088157184, 'peak_bytes_used': 500816}

torch_xla.core.xla_model.get_stablehlo(tensors: Optional[List[Tensor]] = None) → str

以字符串格式获取计算图的 StableHLO。

如果 tensors 不为空，则将转储以 tensors 为输出的图。如果 tensors 为空，则将转储整个计算图。

对于推理图，建议将模型输出传递给 tensors。对于训练图，识别“输出”并不直接。建议使用空的 tensors。

要为 StableHLO 启用源代码行信息，请设置环境变量 XLA_HLO_DEBUG=1。

参数

tensors (list[torch.Tensor], optional) – 代表 StableHLO 图的输出/根的张量。

返回

StableHLO 模块（字符串格式）。

torch_xla.core.xla_model.get_stablehlo_bytecode(tensors: Optional[Tensor] = None) → bytes

以字节码格式获取计算图的 StableHLO。

如果 tensors 不为空，则将转储以 tensors 为输出的图。如果 tensors 为空，则将转储整个计算图。

对于推理图，建议将模型输出传递给 tensors。对于训练图，识别“输出”并不直接。建议使用空的 tensors。

参数

tensors (list[torch.Tensor], optional) – 代表 StableHLO 图的输出/根的张量。

返回

StableHLO 模块（字节码格式）。

distributed

class torch_xla.distributed.parallel_loader.MpDeviceLoader(loader, device, **kwargs)

使用后台数据上传包装现有的 PyTorch DataLoader。

此类应仅与多进程数据并行一起使用。它将用 ParallelLoader 包装传入的 DataLoader，并返回当前设备的 per_device_loader。

参数

• loader (torch.utils.data.DataLoader) – 要包装的 PyTorch DataLoader。

• device (torch.device…) – 数据必须发送到的设备。

• kwargs – ParallelLoader 构造函数的命名参数。

示例

>>> device = torch_xla.device()
>>> train_device_loader = MpDeviceLoader(train_loader, device)

torch_xla.distributed.xla_multiprocessing.spawn(fn, args=(), nprocs=None, join=True, daemon=False, start_method='spawn')

启用多进程复制。

参数

• fn (callable) – 复制中每个进程将调用的函数，该函数将以复制中进程的全局索引作为第一个参数，然后是 args 中传递的参数。

• args (tuple) – fn 的参数。默认值：空元组

• nprocs (python:int) – 复制的进程/设备数量。目前，如果指定，则只能是 1 或 None（将自动转换为最大设备数）。其他数字将导致 ValueError。

• join (bool) – 调用是否应阻塞等待已启动进程的完成。默认值：True

• daemon (bool) – 启动的进程是否应设置 daemon 标志（请参阅 Python 多进程 API）。默认值：False

• start_method (string) – Python multiprocessing 进程创建方法。默认值：spawn

返回

由 torch.multiprocessing.spawn API 返回的相同对象。如果 nprocs 为 1，则将直接调用 fn 函数，并且 API 将返回 None。

spmd

torch_xla.distributed.spmd.mark_sharding(t: Union[Tensor, XLAShardedTensor], mesh: Mesh, partition_spec: Tuple[Optional[Union[Tuple, int, str]], ...]) → XLAShardedTensor

使用 XLA 分区规范注解提供的张量。在内部，它为 XLA SpmdPartitioner 传递注解相应的 XLATensor 以进行分片。

参数

• t (Union[torch.Tensor, XLAShardedTensor]) – 要用 partition_spec 注解的输入张量。

• mesh (Mesh) – 描述逻辑 XLA 设备拓扑和底层设备 ID。

• partition_spec (Tuple[Tuple, python:int, str, None]) – 设备网格维度索引或 None 的元组。每个索引是一个整数，如果网格轴已命名，则为字符串，或整数或字符串的元组。这指定了每个输入秩如何分片（索引到 mesh_shape）或复制（None）。当指定元组时，相应的输入张量轴将沿元组中的所有逻辑轴进行分片。请注意，元组中指定网格轴的顺序将影响最终的分片。

• dynamo_custom_op (bool) – 如果设置为 True，则调用 mark_sharding 的 dynamo 自定义操作变体，使其本身可被 dynamo 识别和跟踪。

示例

>>> import torch_xla.runtime as xr
>>> import torch_xla.distributed.spmd as xs
>>> mesh_shape = (4, 2)
>>> num_devices = xr.global_runtime_device_count()
>>> device_ids = np.array(range(num_devices))
>>> mesh = Mesh(device_ids, mesh_shape, ('x', 'y'))
>>> input = torch.randn(8, 32).to(xm.xla_device())
>>> xs.mark_sharding(input, mesh, (0, None)) # 4-way data parallel
>>> linear = nn.Linear(32, 10).to(xm.xla_device())
>>> xs.mark_sharding(linear.weight, mesh, (None, 1)) # 2-way model parallel

torch_xla.distributed.spmd.clear_sharding(t: Union[Tensor, XLAShardedTensor]) → Tensor

清除输入张量的分片注解，并返回一个 cpu 转换的张量。这是一个就地操作，但也会返回相同的 torch.Tensor。

参数

t (Union[torch.Tensor, XLAShardedTensor]) – 我们想要清除分片的张量

返回

无分片的张量。

返回类型

t (torch.Tensor)

示例

>>> import torch_xla.distributed.spmd as xs
>>> torch_xla.runtime.use_spmd()
>>> t1 = torch.randn(8,8).to(torch_xla.device())
>>> mesh = xs.get_1d_mesh()
>>> xs.mark_sharding(t1, mesh, (0, None))
>>> xs.clear_sharding(t1)

torch_xla.distributed.spmd.set_global_mesh(mesh: Mesh)

设置可用于当前进程的全局网格。

参数

mesh – (Mesh) 将作为全局网格的 Mesh 对象。

示例

>>> import torch_xla.distributed.spmd as xs
>>> mesh = xs.get_1d_mesh("data")
>>> xs.set_global_mesh(mesh)

torch_xla.distributed.spmd.get_global_mesh() → Optional[Mesh]

获取当前进程的全局 Mesh。

返回

(Optional[Mesh]) 如果设置了全局 Mesh，则返回 Mesh 对象，否则返回 None。

返回类型

mesh

示例

>>> import torch_xla.distributed.spmd as xs
>>> xs.get_global_mesh()

torch_xla.distributed.spmd.get_1d_mesh(axis_name: Optional[str] = None) → Mesh

一个辅助函数，用于返回所有设备都在一个维度上的 Mesh。

参数

axis_name – (Optional[str]) 可选字符串，用于表示 Mesh 的轴名称

返回

Mesh 对象

返回类型

Mesh

示例

>>> # This example is assuming 1 TPU v4-8
>>> import torch_xla.distributed.spmd as xs
>>> mesh = xs.get_1d_mesh("data")
>>> print(mesh.mesh_shape)
(4,)
>>> print(mesh.axis_names)
('data',)

class torch_xla.distributed.spmd.Mesh(device_ids: Union[ndarray, List], mesh_shape: Tuple[int, ...], axis_names: Optional[Tuple[str, ...]] = None)

描述逻辑 XLA 设备拓扑 Mesh 和底层资源。

参数

• device_ids (Union[np.ndarray, List]) – 设备（ID）的扁平化列表，顺序自定义。该列表将重塑为 mesh_shape 数组，并使用类 C 索引顺序填充元素。

• mesh_shape (Tuple[python:int, ...]) – 一个整数元组，描述设备 Mesh 的逻辑拓扑形状，每个元素描述相应轴上的设备数量。

• axis_names (Tuple[str, ...]) – 要分配给 devices 参数维度的资源轴名称序列。其长度应与 devices 的秩匹配。

示例

>>> mesh_shape = (4, 2)
>>> num_devices = len(xm.get_xla_supported_devices())
>>> device_ids = np.array(range(num_devices))
>>> mesh = Mesh(device_ids, mesh_shape, ('x', 'y'))
>>> mesh.get_logical_mesh()
>>> array([[0, 1],
          [2, 3],
          [4, 5],
          [6, 7]])
>>> mesh.shape()
OrderedDict([('x', 4), ('y', 2)])

class torch_xla.distributed.spmd.HybridMesh(*, ici_mesh_shape: Tuple[int, ...], dcn_mesh_shape: Optional[Tuple[int, ...]] = None, axis_names: Optional[Tuple[str, ...]] = None)

创建由 ICI 和 DCN 网络连接的设备的混合设备 Mesh。

逻辑 Mesh 的形状应按网络强度递增排序，例如 [replica, data, model]，其中 model 需要最多的网络通信。

参数

• ici_mesh_shape – 内部连接设备的逻辑 Mesh 的形状。

• dcn_mesh_shape – 外部连接设备的逻辑 Mesh 的形状。

示例

>>> # This example is assuming 2 slices of v4-8.
>>> ici_mesh_shape = (1, 4, 1) # (data, fsdp, tensor)
>>> dcn_mesh_shape = (2, 1, 1)
>>> mesh = HybridMesh(ici_mesh_shape, dcn_mesh_shape, ('data','fsdp','tensor'))
>>> print(mesh.shape())
>>> >> OrderedDict([('data', 2), ('fsdp', 4), ('tensor', 1)])

experimental

torch_xla.experimental.eager_mode(enable: bool)

配置 torch_xla 的默认执行模式。

在 eager 模式下，只有通过 `torch_xla.compile` 编译的函数才会被追踪和编译。其他 torch 操作将以 eager 方式执行。

debug

torch_xla.debug.metrics.metrics_report()

检索包含完整指标和计数器报告的字符串。

torch_xla.debug.metrics.short_metrics_report(counter_names: Optional[list] = None, metric_names: Optional[list] = None)

检索包含完整指标和计数器报告的字符串。

参数

• counter_names (list) – 需要打印数据的计数器名称列表。

• metric_names (list) – 需要打印数据的指标名称列表。

torch_xla.debug.metrics.counter_names()

检索所有当前活动的计数器名称。

torch_xla.debug.metrics.counter_value(name)

返回活动计数器的值。

参数

name (string) – 需要检索其值的计数器的名称。

返回

计数器的值，作为整数。

torch_xla.debug.metrics.metric_names()

检索所有当前活动的指标名称。

torch_xla.debug.metrics.metric_data(name)

返回活动指标的数据。

参数

name (string) – 需要检索其数据的指标的名称。

返回

指标数据，是一个元组 (TOTAL_SAMPLES, ACCUMULATOR, SAMPLES)。TOTAL_SAMPLES 是已发布到该指标的样本总数。一个指标只保留给定数量的样本（在循环缓冲区中）。ACCUMULATOR 是 TOTAL_SAMPLES 样本的总和。SAMPLES 是一个 (TIME, VALUE) 元组列表。