torch._utils 的源代码

import copyreg
import functools
import sys
import traceback
import warnings
from collections import defaultdict
from typing import Any, DefaultDict, List, Optional

import torch


def _type(self, dtype=None, non_blocking=False, **kwargs):
    """如果未提供 `dtype`，则返回类型，否则将此对象转换为指定类型。
    如果已经是正确类型，则不执行复制并返回原始对象。

    参数:
        dtype (类型或字符串): 所需类型
        non_blocking (bool): 如果为 ``True``，并且源在固定内存中
            且目标在 GPU 上或反之，则复制将相对于主机异步执行。否则，该参数无效。
        **kwargs: 为了兼容性，可能包含 ``async`` 键，代替 ``non_blocking`` 参数。``async`` 参数已弃用。
    """
    non_blocking = _get_async_or_non_blocking("type", non_blocking, kwargs)
    if dtype is None:
        return self.__module__ + "." + self.__class__.__name__

    if isinstance(dtype, str):
        dtype = _import_dotted_name(dtype)
    if dtype == type(self):
        return self
    if self.is_sparse:
        if not dtype.is_sparse:
            raise RuntimeError("无法将稀疏张量转换为稠密张量")
        new_module_name = dtype.__module__.replace(".sparse", "")
        new_values_type_name = new_module_name + "." + dtype.__name__
        new_values = torch.Tensor._values(self).type(new_values_type_name, non_blocking)
        new_indices_type_name = new_module_name + ".LongTensor"
        new_indices = torch.Tensor._indices(self).type(
            new_indices_type_name, non_blocking
        )
        return dtype(new_indices, new_values, self.size())
    if dtype.is_sparse:
        raise RuntimeError("无法将稠密张量转换为稀疏张量")
    return dtype(self.size()).copy_(self, non_blocking)


def _hpu(self, device=None, non_blocking=False, **kwargs):
    """返回此对象在 HPU 内存中的副本。

    如果此对象已经在 HPU 内存中且在正确的设备上，则不执行复制并返回原始对象。

    参数:
        device (int): 目标 HPU 设备 ID。默认为当前设备。
        non_blocking (bool): 如果为 ``True`` 且源在固定内存中，
            则复制将相对于主机异步执行。否则，该参数无效。
        **kwargs: 为了兼容性，可能包含 ``async`` 键，代替 ``non_blocking`` 参数。
    """
    non_blocking = _get_async_or_non_blocking("hpu", non_blocking, kwargs)
    hpu = getattr(torch, "hpu", None)
    assert hpu is not None, "HPU 设备模块未加载"
    if self.is_hpu:
        if device is None:
            device = hpu.current_device()
        if self.get_device() == device:
            return self
    else:
        if device is None:
            device = -1
    with hpu.device(device):
        assert not self.is_sparse, "HPU 张量不支持稀疏存储"
        untyped_storage = torch.UntypedStorage(self.size(), device=torch.device("hpu"))
        untyped_storage.copy_(self, non_blocking)
        return untyped_storage


def _cuda(self, device=None, non_blocking=False, **kwargs):
    """返回此对象在 CUDA 内存中的副本。

    如果此对象已经在 CUDA 内存中且在正确的设备上，则不执行复制并返回原始对象。

    参数:
        device (int): 目标 GPU ID。默认为当前设备。
        non_blocking (bool): 如果为 ``True`` 且源在固定内存中，
            则复制将相对于主机异步执行。否则，该参数无效。
        **kwargs: 为了兼容性，可能包含 ``async`` 键，代替 ``non_blocking`` 参数。
    """
    non_blocking = _get_async_or_non_blocking("cuda", non_blocking, kwargs)
    if self.is_cuda:
        if device</