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 worker ID 具有任何可预测的关系，也不保证它们在每个主机上是连续的。

torch_xla.runtime.local_ordinal() → int

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

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

torch_xla.runtime.get_master_ip() → str

检索运行时的主 worker IP。这将调用特定于后端的发现 API。

返回

主 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) – 此 worker 是否应具有对缓存的写入权限。

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]]

对输入张量执行就地 reduce 操作。

参数

• 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) – reduce 后应用的默认缩放值。默认值：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 has a mix of layout constrained”的错误消息时，请取消固定布局。

返回

如果传递单个 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) – gather 维度。默认值：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 has a mix of layout constrained”的错误消息时，请取消固定布局。

返回

一个张量，在 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 has a mix of layout constrained”的错误消息时，请取消固定布局。

返回

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]

等待所有网格客户端到达指定的 rendezvous。

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

参数

• tag (string) – 要加入的 rendezvous 的名称。

• payload (bytes, optional) – 要发送到 rendezvous 的有效负载。

• replicas (list, python:int) – 参与 rendezvous 的副本序号。空列表表示网格中的所有副本。默认值：[]

返回

所有其他核心交换的有效负载，核心序号为 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]

执行图外客户端网格缩减。

参数

• tag (string) – 要加入的 rendezvous 的名称。

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

• 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 包装传入的数据加载器，并为当前设备返回 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 分区规范注释提供的张量。在内部，它将相应的 XLATensor 注释为针对 XLA SpmdPartitioner 传递进行分片。

参数

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

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

• partition_spec (Tuple[Tuple, python:int, str, None]) – 设备网格维度索引或 None 的元组。每个索引都是 int、str（如果网格轴已命名）或 int 或 str 的元组。这指定了每个输入秩如何分片（索引到 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]

获取当前进程的全局网格。

返回

(Optional[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

辅助函数，用于返回网格，其中所有设备都在一个维度中。

参数

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

返回

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 设备拓扑网格和底层资源。

参数

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

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

• 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 网络连接的设备的混合设备网格。

逻辑网格的形状应按网络密集程度递增排序，例如 [副本, 数据, 模型]，其中 mdl 具有最多的网络通信需求。

参数

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

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

示例

>>> # 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) 元组的列表。