# Source code for tensorplay Source: https://www.tensorplay.cn/docs/_modules/tensorplay.html ``` """ The tensorplay package offers a simple deep-learning framework designed for educational purposes and small-scale experiments. It defines a data structure for multidimensional arrays called Tensor, on which it encapsulates mathematical operations. It has a CUDA counterpart, that enables you to run your tensor computations on an NVIDIA GPU with compute capability >= 3.0. """ import builtins import ctypes from enum import IntEnum from math import e, inf, nan, pi # noqa: F401 import glob import importlib import inspect import os import sys import types import threading from typing import ( Any as _Any, TYPE_CHECKING, ) from typing_extensions import TypeIs as _TypeIs # Generated at build time by tools/generate_tensorplay_version.py (CMake # generated version module. from tensorplay.version import __version__ as __version__ # The OpenMP runtime parks its worker team only after a fixed spin-out, which # defaults to several milliseconds of full-core spinning per worker. Op # sequences reach the next region far sooner than that, so a running process # keeps every worker core hot even when each region is tiny. A ~100us window # still spans back-to-back regions (measured wall-clock identical) while # letting idle workers sleep between bursts. Explicit user settings win. if "GOMP_SPINCOUNT" not in os.environ and "OMP_WAIT_POLICY" not in os.environ: os.environ["GOMP_SPINCOUNT"] = "10000" # ------------------------------------------------------------------------- # DLL Loading (Windows) # ------------------------------------------------------------------------- if sys.platform == 'win32': def _load_dll_libraries(): # Adapted from TensorPlay import sysconfig # Helper to add DLL directory safely def _add_dll_directory(path): if os.path.exists(path): try: os.add_dll_directory(path) except (OSError, AttributeError): pass # 1. Package's own lib directory (p10.dll, tpx.dll, stax.dll, dnnl.dll, etc.) package_lib_path = os.path.join(os.path.dirname(__file__), 'lib') _add_dll_directory(package_lib_path) # Add to PATH immediately as fallback os.environ["PATH"] = package_lib_path + ";" + os.environ["PATH"] # 2. Package's root directory (sometimes DLLs are here) _add_dll_directory(os.path.dirname(__file__)) # 3. Conda/Python Library/bin (for MKL, OneDNN, etc.) py_dll_path = os.path.join(sys.exec_prefix, 'Library', 'bin') _add_dll_directory(py_dll_path) # 4. VirtualEnv support if sys.exec_prefix != sys.base_exec_prefix: base_py_dll_path = os.path.join(sys.base_exec_prefix, "Library", "bin") _add_dll_directory(base_py_dll_path) # 5. User site-packages Library/bin userbase = sysconfig.get_config_var('userbase') if userbase: user_dll_path = os.path.join(userbase, 'Library', 'bin') _add_dll_directory(user_dll_path) # 6. Explicitly load DLLs to ensure dependencies are resolved kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) with_load_library_flags = hasattr(kernel32, "AddDllDirectory") prev_error_mode = kernel32.SetErrorMode(0x0001) # Pre-load critical dependencies in order package_lib_path = os.path.join(os.path.dirname(__file__), 'lib') # Order matters! Dependencies first. # MKL -> CUDA/cuDNN -> p10 -> tpx -> stax # 1. Pre-load MKL (if present) mkl_dlls = glob.glob(os.path.join(package_lib_path, "mkl_*.dll")) # Load mkl_core and mkl_sequential first if they exist priority_mkl = ["mkl_core", "mkl_sequential", "mkl_intel_lp64", "mkl_def", "mkl_avx2"] sorted_mkl = [] for name in priority_mkl: for dll in mkl_dlls: if name in os.path.basename(dll): sorted_mkl.append(dll) # Add remaining MKL DLLs for dll in mkl_dlls: if dll not in sorted_mkl: sorted_mkl.append(dll) # 2. Pre-load CUDA/cuDNN (if present) cuda_dlls = glob.glob(os.path.join(package_lib_path, "cudart*.dll")) + \ glob.glob(os.path.join(package_lib_path, "cublas*.dll")) + \ glob.glob(os.path.join(package_lib_path, "cudnn*.dll")) + \ glob.glob(os.path.join(package_lib_path, "curand*.dll")) # 3. Core Libraries core_dlls = [ os.path.join(package_lib_path, "p10.dll"), os.path.join(package_lib_path, "tpx.dll"), os.path.join(package_lib_path, "stax.dll") ] all_dlls = sorted_mkl + cuda_dlls + core_dlls path_patched = False for dll in all_dlls: if not os.path.exists(dll): continue if "OpenCL" in dll: continue is_loaded = False if with_load_library_flags: # LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR res = kernel32.LoadLibraryExW(dll, None, 0x00001100) if res: is_loaded = True if not is_loaded: # Fallback if not path_patched: os.environ["PATH"] = package_lib_path + ";" + os.environ["PATH"] path_patched = True res = kernel32.LoadLibraryW(dll) kernel32.SetErrorMode(prev_error_mode) _load_dll_libraries() del _load_dll_libraries # ------------------------------------------------------------------------- # ------------------------------------------------------------------------- # preloads component CUDA wheels when the extension reports a missing SONAME. # Keep the same behavior here: a normal CPU-only import does not eagerly load # CUDA, while a CUDA build remains importable from NVIDIA Python wheels. elif sys.platform.startswith('linux'): # The CUDA wheels consolidate their libraries under a per-toolkit-major # directory (nvidia/cuNN, e.g. cu13 for the 13.x generation). Derive the # directory from the toolkit version baked into the build metadata so the # scan follows the wheel generation this package was linked against. A # package whose metadata names no toolkit (CPU build) skips the scan; a # preloaded object carrying a foreign SONAME satisfies nothing and the # extension retry surfaces the honest missing-soname error. try: from .version import cuda as _tp_cuda_toolkit except ImportError: _tp_cuda_toolkit = None _tp_cuda_wheel_dir = _tp_cuda_toolkit.split('.', 1)[0] if _tp_cuda_toolkit else None def _get_cuda_dep_paths(path, lib_folder, lib_name): paths = [] # Exact wheel layout only (nvidia/cuda_runtime, nvidia/cublas, ...). # A broad 'nvidia/cu*' glob would also swallow foreign major-version # wheels like nvidia/cu13 and dlopen the wrong libcudart SONAME, paths.extend(glob.glob(os.path.join(path, 'nvidia', lib_folder, 'lib', lib_name))) paths.extend(glob.glob(os.path.join(path, lib_folder, 'lib', lib_name))) if _tp_cuda_wheel_dir: paths.extend(glob.glob(os.path.join( path, 'nvidia', 'cu' + _tp_cuda_wheel_dir, 'lib', lib_name))) if not paths and '.so.' in lib_name: stem = lib_name.split('.so.', 1)[0] paths.extend(glob.glob(os.path.join(path, 'nvidia', lib_folder, 'lib', stem + '.so'))) paths.extend(glob.glob(os.path.join(path, lib_folder, 'lib', stem + '.so'))) if _tp_cuda_wheel_dir: paths.extend(glob.glob(os.path.join( path, 'nvidia', 'cu' + _tp_cuda_wheel_dir, 'lib', stem + '.so'))) return paths def _preload_cuda_lib(lib_folder, lib_name, required=True): for path in sys.path: candidates = _get_cuda_dep_paths(path, lib_folder, lib_name) if candidates: ctypes.CDLL(candidates[0]) return package_lib_path = os.path.join(os.path.dirname(__file__), 'lib') candidates = glob.glob(os.path.join(package_lib_path, lib_name)) if not candidates and '.so.' in lib_name: stem = lib_name.split('.so.', 1)[0] candidates = glob.glob(os.path.join(package_lib_path, stem + '.so')) if candidates: ctypes.CDLL(candidates[0]) return if required: raise ImportError(f'{lib_name} not found in TensorPlay CUDA dependency paths') def _preload_cuda_deps(err=None): if err is not None: message = str(err) # The package later rebinds any() to the tensor reduction. cuda_error = builtins.any( token in message for token in ('libcud', 'libcublas', 'libcurand', 'libnvrtc', 'libnvjpeg') ) if not cuda_error: raise err # a mixed CUDA installation through cublas' transitive RUNPATH. for lib_folder, lib_name, required in ( ('cublas', 'libcublasLt.so.*[0-9]', True), ('cublas', 'libcublas.so.*[0-9]', True), ('cudnn', 'libcudnn.so.*[0-9]', True), ('cuda_nvrtc', 'libnvrtc.so.*[0-9]', False), ('cuda_runtime', 'libcudart.so.*[0-9]', True), ('curand', 'libcurand.so.*[0-9]', True), # The image IO bindings link nvjpeg; a toolkit without it # surfaces the honest missing-soname error on the retry. ('nvjpeg', 'libnvjpeg.so.*[0-9]', False), ): _preload_cuda_lib(lib_folder, lib_name, required=required) # ------------------------------------------------------------------------- # Worker wait # ------------------------------------------------------------------------- # How long an idle worker thread keeps polling for the next parallel region # before it goes to sleep is read by the threading runtime once, when it is # loaded. Left unsaid, some builds of the runtime sleep at once: every region # then starts by waking its workers, and a training step made of a thousand # short regions spends a fifth of its time doing that. A few milliseconds of # polling is the runtime's long-standing default, so it is asked for by name # unless the environment already says how workers should wait. if "GOMP_SPINCOUNT" not in os.environ and "OMP_WAIT_POLICY" not in os.environ: os.environ["GOMP_SPINCOUNT"] = "300000" # ------------------------------------------------------------------------- # Core Imports # ------------------------------------------------------------------------- try: from . import _C as _C except (ImportError, OSError) as _load_error: if sys.platform.startswith('linux') and '_preload_cuda_deps' in globals(): _preload_cuda_deps(_load_error) from . import _C as _C else: raise if not hasattr(_C, "_log_api_usage_once"): _C._log_api_usage_once = lambda *args, **kwargs: None from ._tensor import Tensor from ._C import (tensor, DType, Size, Scalar, SymInt, SymBool, SymFloat, sym_float, sym_int, sym_not, sym_min, sym_max, sym_ite, sym_sum, Device, DeviceType, from_dlpack, to_dlpack, from_numpy, frombuffer, set_printoptions, default_generator, manual_seed, seed, initial_seed, Generator, UntypedStorage, is_storage, get_rng_state, set_rng_state, set_num_threads, get_num_threads, get_thread_num, in_parallel_region, get_parallel_info) from .autograd import (no_grad, enable_grad, set_grad_enabled, is_grad_enabled, inference_mode, is_inference_mode_enabled) if hasattr(_C, "get_num_interop_threads"): from ._C import get_num_interop_threads, set_num_interop_threads else: # The compiled extension predates the interop thread entry points; expose # wrappers over the shared thread pool until the extension is rebuilt. def get_num_interop_threads(): """Returns the number of threads used for interop parallelism on CPU. The parallel layer runs a single shared thread pool, so this equals the number of intraop threads. """ return get_num_threads() def set_num_interop_threads(nthreads): """Sets the number of threads used for interop parallelism on CPU. The parallel layer runs a single shared thread pool, so the request is accepted only when it matches the pool size configured via ``set_num_threads``; any other size raises RuntimeError. """ if nthreads <= 0: raise ValueError("Expected positive number of threads") current = get_num_threads() if nthreads != current: raise RuntimeError( f"cannot set number of interop threads to {nthreads}: the " f"parallel layer uses a single shared thread pool whose size " f"({current}) is configured via set_num_threads before " f"parallel work has started") from .serialization import save, load, inspect_checkpoint from .serialization import archive as _serialization_archive from .random import fork_rng, thread_safe_generator # ------------------------------------------------------------------------- # DType Aliases # ------------------------------------------------------------------------- device = Device dtype = DType uint8 = DType.uint8 int8 = DType.int8 int16 = DType.int16 uint16 = DType.uint16 uint32 = DType.uint32 uint64 = DType.uint64 int32 = DType.int32 int64 = DType.int64 float32 = DType.float32 float64 = DType.float64 float16 = DType.float16 bfloat16 = DType.bfloat16 complex32 = DType.complex32 complex64 = DType.complex64 complex128 = DType.complex128 bcomplex32 = DType.bcomplex32 bool = DType.bool undefined = DType.undefined qint8 = DType.qint8 quint8 = DType.quint8 qint32 = DType.qint32 quint2x4 = DType.quint2x4 quint4x2 = DType.quint4x2 quint1x8 = DType.quint1x8 float8_e4m3fn = DType.float8_e4m3fn float8_e4m3fnuz = DType.float8_e4m3fnuz float8_e5m2 = DType.float8_e5m2 float8_e5m2fnuz = DType.float8_e5m2fnuz float8_e8m0fnu = DType.float8_e8m0fnu half = DType.float16 float = DType.float32 double = DType.float64 short = DType.int16 int = DType.int32 long = DType.int64 cfloat = DType.complex64 cdouble = DType.complex128 chalf = DType.complex32 [docs] class MemoryFormat(IntEnum): CONTIGUOUS = 0 PRESERVE = 1 CHANNELS_LAST = 2 CHANNELS_LAST_3D = 3 contiguous_format = 0 preserve_format = 1 channels_last = 2 channels_last_3d = 3 # Printed as the public module attribute, so generated graph code can # spell the value back. def __repr__(self): return f"tensorplay.{('contiguous_format', 'preserve_format', 'channels_last', 'channels_last_3d')[self.value]}" __str__ = __repr__ [docs] class QScheme(IntEnum): """Quantization scheme attached to a quantized tensor's quantizer.""" PER_TENSOR_AFFINE = 0 PER_CHANNEL_AFFINE = 1 PER_TENSOR_SYMMETRIC = 2 PER_CHANNEL_SYMMETRIC = 3 PER_CHANNEL_AFFINE_FLOAT_QPARAMS = 4 per_tensor_affine = PER_TENSOR_AFFINE per_channel_affine = PER_CHANNEL_AFFINE per_tensor_symmetric = PER_TENSOR_SYMMETRIC per_channel_symmetric = PER_CHANNEL_SYMMETRIC per_channel_affine_float_qparams = PER_CHANNEL_AFFINE_FLOAT_QPARAMS def __repr__(self): return f"tensorplay.{self.name.lower()}" __str__ = __repr__ per_tensor_affine = QScheme.PER_TENSOR_AFFINE per_channel_affine = QScheme.PER_CHANNEL_AFFINE per_tensor_symmetric = QScheme.PER_TENSOR_SYMMETRIC per_channel_symmetric = QScheme.PER_CHANNEL_SYMMETRIC per_channel_affine_float_qparams = QScheme.PER_CHANNEL_AFFINE_FLOAT_QPARAMS [docs] class Layout(IntEnum): """ ``tensor.layout`` returns one of these values; dense tensors report STRIDED. """ SPARSE_COO = 0 SPARSE_CSR = 1 SPARSE_CSC = 2 SPARSE_BSR = 3 SPARSE_BSC = 4 STRIDED = 5 sparse_coo = 0 sparse_csr = 1 sparse_csc = 2 sparse_bsr = 3 sparse_bsc = 4 strided = 5 def __repr__(self): return f"tensorplay.{self.name.lower()}" __str__ = __repr__ sparse_coo = Layout.SPARSE_COO sparse_csr = Layout.SPARSE_CSR sparse_csc = Layout.SPARSE_CSC sparse_bsr = Layout.SPARSE_BSR sparse_bsc = Layout.SPARSE_BSC strided = Layout.STRIDED contiguous_format = MemoryFormat.CONTIGUOUS preserve_format = MemoryFormat.PRESERVE channels_last = MemoryFormat.CHANNELS_LAST channels_last_3d = MemoryFormat.CHANNELS_LAST_3D # The enum class under its lowercase alias, plus the historical name for the # contiguous format kept for source compatibility. memory_format = MemoryFormat # Link-time capabilities of the compiled backend. has_lapack = True has_spectral = True compiled_with_cxx11_abi = True def getArtifactLogger(name: str, artifact: str): """A logger for one named artifact, separate from the module's own. A piece of work that leaves something behind -- a trace, a set of choices, a schedule -- is worth being able to look at on its own, without the rest of what the module says. So a logger is named for the artifact rather than for the place, and can be turned on and off without turning anything else on or off with it. """ import logging return logging.getLogger(f"{name}.{artifact}") __all__ = [ "Tensor", "tensor", "from_dlpack", "Scalar", "SymInt", "SymBool", "SymFloat", "sym_float", "sym_int", "sym_not", "sym_min", "sym_max", "sym_ite", "sym_sum", "DeviceType", "device", "dtype", "Size", "MemoryFormat", "contiguous_format", "preserve_format", "channels_last", "channels_last_3d", "memory_format", "has_lapack", "has_spectral", "compiled_with_cxx11_abi", "Layout", "sparse_coo", "sparse_csr", "sparse_csc", "sparse_bsr", "sparse_bsc", "strided", "uint8", "int8", "int16", "uint16", "uint32", "uint64", "int32", "int64", "float16", "bfloat16", "float32", "float64", "complex32", "complex64", "complex128", "bcomplex32", "bool", "qint8", "quint8", "qint32", "quint2x4", "quint4x2", "quint1x8", "QScheme", "per_tensor_affine", "per_channel_affine", "per_tensor_symmetric", "per_channel_symmetric", "per_channel_affine_float_qparams", "half", "float", "double", "short", "int", "long", "cfloat", "cdouble", "chalf", "save", "load", "inspect_checkpoint", "as_tensor", "quantized_lstm", "quantized_gru", "no_grad", "enable_grad", "set_grad_enabled", "is_grad_enabled", "is_inference_mode_enabled", "allclose", "index", "compile", "compiler", "graph", "library", "set_num_threads", "get_num_threads", "get_thread_num", "set_num_interop_threads", "get_num_interop_threads", "in_parallel_region", "get_parallel_info", "default_generator", "manual_seed", "seed", "initial_seed", "Generator", "UntypedStorage", "is_storage", "get_rng_state", "set_rng_state", "fork_rng", "thread_safe_generator", "DeviceMismatchError", "cond", "while_loop", "getArtifactLogger", "__config__", ] # Append functional API to __all__ __all__.extend([ "abs", "acos", "acosh", "adaptive_avg_pool2d", "adaptive_avg_pool2d_backward", "adaptive_max_pool2d", "adaptive_max_pool2d_backward", "all", "angle", "any", "arange", "argmax", "argmin", "asin", "asinh", "atan", "atan2", "atanh", "avg_pool2d", "avg_pool2d_backward", "batch_norm", "batch_norm_backward", "bernoulli", "cat", "ceil", "chunk", "clamp", "clamp_backward", "constant_pad_nd", "constant_pad_nd_backward", "conv1d", "conv1d_grad_bias", "conv1d_grad_input", "conv1d_grad_weight", "conv2d", "conv2d_grad_bias", "conv2d_grad_input", "conv2d_grad_weight", "conv3d", "conv3d_grad_bias", "conv3d_grad_input", "conv3d_grad_weight", "conv_transpose2d", "conv_transpose2d_grad_bias", "conv_transpose2d_grad_input", "conv_transpose2d_grad_weight", "conv_transpose3d", "conv_transpose3d_grad_bias", "conv_transpose3d_grad_input", "conv_transpose3d_grad_weight", "cos", "cosh", "embedding", "embedding_dense_backward", "empty", "empty_like", "eq", "exp", "eye", "floor", "full", "full_like", "ge", "gelu", "group_norm", "group_norm_backward", "gt", "instance_norm", "instance_norm_backward", "layer_norm", "layer_norm_backward", "le", "lerp", "linspace", "log", "log_softmax", "logspace", "lt", "masked_select", "matmul", "max", "max_pool2d", "max_pool2d_backward", "mean", "median", "min", "mm", "mse_loss", "mse_loss_backward", "ne", "neg", "nll_loss", "nll_loss_backward", "norm", "normal", "ones", "ones_like", "permute", "permute_backward", "poisson", "pow", "prod", "rand", "rand_like", "randint", "randint_like", "randn", "randn_like", "randperm", "relu", "reshape", "round", "rsqrt", "sigmoid", "sign", "silu", "sin", "sinh", "softmax", "split", "sqrt", "square", "sparse_compressed_tensor", "sparse_csr_tensor", "sparse_csc_tensor", "sparse_bsr_tensor", "sparse_bsc_tensor", "squeeze", "squeeze_backward", "stack", "std", "sum", "t", "tan", "tanh", "threshold_backward", "transpose", "unbind", "unsqueeze", "var", "zeros", "zeros_like", "where", "maximum", "minimum", "addcmul", "addcdiv", "view_as_real", "view_as_complex", "is_complex", ]) # Please keep this list sorted # assert __all__ == sorted(__all__) import functools from ._C import ( abs, acos, acosh, adaptive_avg_pool2d, adaptive_avg_pool2d_backward, adaptive_max_pool2d, adaptive_max_pool2d_backward, all, angle, any, arange, argmax, argmin, asin, asinh, atan, atan2, atanh, avg_pool2d, avg_pool2d_backward, batch_norm, batch_norm_backward, bernoulli, cat, ceil, chunk, clamp, clamp_backward, constant_pad_nd, constant_pad_nd_backward, conv1d, conv1d_grad_bias, conv1d_grad_input, conv1d_grad_weight, conv2d, conv2d_grad_bias, conv2d_grad_input, conv2d_grad_weight, conv3d, conv3d_grad_bias, conv3d_grad_input, conv3d_grad_weight, conv_transpose2d, conv_transpose2d_grad_bias, conv_transpose2d_grad_input, conv_transpose2d_grad_weight, conv_transpose3d, conv_transpose3d_grad_bias, conv_transpose3d_grad_input, conv_transpose3d_grad_weight, cos, cosh, embedding, embedding_dense_backward, empty, empty_like, eq, exp, eye, floor, full, full_like, ge, gelu, group_norm, group_norm_backward, gt, instance_norm, instance_norm_backward, layer_norm, layer_norm_backward, le, lerp, linspace, log, log_softmax, logspace, lt, masked_select, matmul, max, max_pool2d, max_pool2d_backward, mean, median, min, mm, mse_loss, mse_loss_backward, ne, neg, nll_loss, nll_loss_backward, norm, normal, ones, ones_like, permute, permute_backward, poisson, pow, prod, rand, rand_like, randint, randint_like, randn, randn_like, randperm, relu, reshape, round, rsqrt, sigmoid, sign, silu, sin, sinh, softmax, split, sqrt, square, squeeze, squeeze_backward, stack, std, sum, t, tan, tanh, threshold_backward, transpose, unbind, unsqueeze, var, zeros, zeros_like, as_tensor, ) # Catchable device-mismatch error raised by the C++ core (RuntimeError # subclass); kept explicit because it is an exception type, not an op. DeviceMismatchError = _C.DeviceMismatchError __attr_name, __obj = "", None for __attr_name in dir(_C): if __attr_name[0] != "_" and not __attr_name.endswith("Base"): __all__.append(__attr_name) __obj = getattr(_C, __attr_name) if callable(__obj) or inspect.isclass(__obj): if os.getenv("TENSORPLAY_BUILDING_STUBS") == "1": continue if __obj.__module__ != __name__: try: __obj.__module__ = __name__ except AttributeError: # Fallback: wrap it if it's a function and module wasn't updated # (pybind11 functions are builtin_function_or_method, which # rejects __module__ writes) if callable(__obj): def make_wrapper(f): @functools.wraps(f) def wrapper(*args, **kwargs): return f(*args, **kwargs) wrapper.__module__ = __name__ return wrapper wrapper = make_wrapper(__obj) # Overwrite a name already installed by the native binding. if __attr_name in globals(): globals()[__attr_name] = wrapper if __attr_name not in globals(): globals()[__attr_name] = __obj elif __attr_name == "TensorBase": if hasattr(sys.modules[__name__], __attr_name): delattr(sys.modules[__name__], __attr_name) del __attr_name, __obj if not TYPE_CHECKING: def _import_extension_to_sys_modules(module, memo=None): """ Recursively import submodules of a C extension module into sys.modules. """ if memo is None: memo = set() if module in memo: return memo.add(module) module_name = module.__name__ for name in dir(module): member = getattr(module, name) member_name = getattr(member, "__name__", "") if inspect.ismodule(member) and member_name.startswith(module_name): sys.modules.setdefault(member_name, member) _import_extension_to_sys_modules(member, memo) _import_extension_to_sys_modules(_C) del _import_extension_to_sys_modules from .functional import * from .functional import _assert_async from .functional import _unsafe_index_put def unique(input, sorted=True, return_inverse=False, return_counts=False): """Return unique values and optionally inverse indices and counts. The native op always computes all three outputs; this wrapper selects the requested ``return_inverse`` / ``return_counts`` values. """ values, inverse, counts = _C.unique(input, sorted, True, True) if return_inverse and return_counts: return values, inverse, counts if return_counts: return values, counts if return_inverse: return values, inverse return values def _assert(condition: object, message: object) -> None: r"""Python's ``assert`` as a call, so graph capture can record it.""" if type(condition) is not Tensor and overrides.has_tensorplay_function( (condition,) ): return overrides.handle_tensorplay_function( _assert, (condition,), condition, message ) if not condition: raise AssertionError(message) from ._shape_funcs import * from ._composite_funcs import * from ._einsum import einsum from .utils.comparison import allclose from ._ops import ops as ops def _register_checks() -> None: """Keep checks in captured graphs although nothing reads their result.""" from . import functional from .graph.node import _side_effectful_targets native = ops.tp _side_effectful_targets.update({ _assert, functional._assert_async, functional._assert_scalar, functional._assert_tensor_metadata, functional.sym_constrain_range, functional.sym_constrain_range_for_size, native._assert_async.default, native._assert_async.msg, native._assert_scalar.default, native._assert_tensor_metadata.default, native.sym_constrain_range.default, native.sym_constrain_range_for_size.default, }) _register_checks() del _register_checks from . import _stax from .compiler._core import compile from . import compiler from . import library from . import profiler # ------------------------------------------------------------------------- # Submodules # ------------------------------------------------------------------------- from . import cuda # Imported after the op names, nn, and cuda are bound: the multiprocessing # reducers pull in tensorplay.primitives (via nn) which references eager op # names, so this must not run during partial initialization. from . import multiprocessing from . import stax from . import backends # Import nn before optim because optim.swa_utils pulls in tensorplay.nn, # which must be fully initialized by then to avoid a partial-import cycle. from . import nn from . import optim from . import graph from . import autograd from . import distributed from . import utils from . import __config__ from . import amp from .amp.autocast_mode import ( # noqa: E402 autocast_decrement_nesting, autocast_increment_nesting, clear_autocast_cache, get_autocast_cpu_dtype, get_autocast_dtype, get_autocast_gpu_dtype, is_autocast_available, is_autocast_cache_enabled, is_autocast_enabled, set_autocast_cache_enabled, set_autocast_dtype, set_autocast_enabled, ) from .amp import ( # noqa: E402 GradScaler, autocast, custom_bwd, custom_fwd, ) __all__.extend([ "GradScaler", "amp", "cuda", "autocast", "autocast_decrement_nesting", "autocast_increment_nesting", "clear_autocast_cache", "custom_bwd", "custom_fwd", "get_autocast_cpu_dtype", "get_autocast_dtype", "get_autocast_gpu_dtype", "is_autocast_available", "is_autocast_cache_enabled", "is_autocast_enabled", "set_autocast_cache_enabled", "set_autocast_dtype", "set_autocast_enabled", ]) [docs] def typename(obj: _Any, /) -> str: """ String representation of the type of an object. This function returns a fully qualified string representation of an object's type. Args: obj (object): The object whose type to represent Returns: str: the type of the object `o` Example: >>> x = tensorplay.tensor([1, 2, 3]) >>> tensorplay.typename(x) 'tensorplay.LongTensor' >>> tensorplay.typename(tensorplay.nn.Parameter) 'tensorplay.nn.parameter.Parameter' """ if isinstance(obj, tensorplay.Tensor): return obj.type() module = getattr(obj, "__module__", "") or "" qualname = "" if hasattr(obj, "__qualname__"): qualname = obj.__qualname__ elif hasattr(obj, "__name__"): qualname = obj.__name__ else: module = obj.__class__.__module__ or "" qualname = obj.__class__.__qualname__ if module in {"", "builtins"}: return qualname return f"{module}.{qualname}" [docs] def is_tensor(obj: _Any, /) -> _TypeIs["tensorplay.Tensor"]: r"""Returns True if `obj` is a TensorPlay tensor. Note that this function is simply doing ``isinstance(obj, Tensor)``. Using that ``isinstance`` check is better for type checking with mypy, and more explicit - so it's recommended to use that instead of ``is_tensor``. Args: obj (object): Object to test Example:: >>> x = tensorplay.tensor([1, 2, 3]) >>> tensorplay.is_tensor(x) True """ return isinstance(obj, tensorplay.Tensor) [docs] def as_tensor(data, dtype=None, device=None): r"""Convert ``data`` into a tensor, sharing storage when possible. If ``data`` is already a tensor with the requested dtype and device, it is returned as-is (no copy). Otherwise it is converted to the requested dtype and device. Args: data (tensor, list, or scalar): Initial data for the tensor. dtype (:class:`tensorplay.DType`, optional): the desired data type of the returned tensor. device (:class:`tensorplay.Device`, str, optional): the device of the constructed tensor. Example:: >>> x = tensorplay.tensor([1.0, 2.0]) >>> tensorplay.as_tensor(x) is x True >>> tensorplay.as_tensor([0, 1, 2], dtype=tensorplay.int64).dtype == tensorplay.int64 True """ if isinstance(data, tensorplay.Tensor): if dtype is not None and data.dtype != dtype: data = data.to(dtype) if device is not None: if not isinstance(device, tensorplay.device): device = tensorplay.device(device) if data.device != device: data = data.to(device) return data return tensor(data, dtype=dtype, device=device) _GLOBAL_DEVICE_CONTEXT = threading.local() # thread-local mode stack semantics); _GLOBAL_DEVICE_CONTEXT is kept for API [docs] def get_default_device() -> "tensorplay.device": r"""Gets the default ``tensorplay.Tensor`` to be allocated on ``device``""" return _C.get_default_device() [docs] def set_default_device(device: Device) -> None: """Sets the default ``tensorplay.Tensor`` to be allocated on ``device``. This does not affect factory function calls which are called with an explicit ``device`` argument. Factory calls will be performed as if they were passed ``device`` as an argument. To only temporarily change the default device instead of setting it globally, use ``with tensorplay.device(device):`` instead. The default device is initially ``cpu``. If you set the default tensor device to another device (e.g., ``cuda``) without a device index, tensors will be allocated on whatever the current device for the device type, even after :func:`tensorplay.cuda.set_device` is called. .. warning:: This function imposes a slight performance cost on every Python call to the tensorplay API (not just factory functions). .. note:: This doesn't affect functions that create tensors that share the same memory as the input, like: :func:`tensorplay.from_numpy` and :func:`tensorplay.frombuffer` Args: device (:class:`tensorplay.device`, str, int, or None): the device to set as default, or ``None`` to clear the override. An integer is interpreted as an index for the current accelerator. Example:: >>> # xdoctest: +SKIP("requires cuda, changes global state") >>> tensorplay.get_default_device() device(type='cpu') >>> tensorplay.set_default_device('cuda') # current device is 0 >>> tensorplay.get_default_device() device(type='cuda', index=0) >>> tensorplay.set_default_device('cuda') >>> tensorplay.cuda.set_device('cuda:1') # current device is 1 >>> tensorplay.get_default_device() device(type='cuda', index=1) >>> tensorplay.set_default_device('cuda:1') >>> tensorplay.get_default_device() device(type='cuda', index=1) """ if isinstance(device, str): device = Device(device) elif isinstance(device, builtins.int): # An integer is interpreted as an index for the current accelerator # (CUDA is TensorPlay's accelerator device type). device = Device(DeviceType.CUDA, device) _C._set_default_device(device) _default_dtype: DType = float32 [docs] def get_default_dtype() -> DType: """Returns the current default floating point dtype (float32 initially, changed by :func:`set_default_dtype`).""" return _default_dtype [docs] def set_default_dtype(d: DType, /) -> None: r""" Sets the default floating point dtype to :attr:`d`. Supports floating point dtype as inputs. Other dtypes will cause tensorplay to raise an exception. When TensorPlay is initialized its default floating point dtype is float32, and the intent of set_default_dtype(float64) is to facilitate NumPy-like type inference. The default floating point dtype is used to: 1. Implicitly determine the default complex dtype. When the default floating type is float16, the default complex dtype is complex32. For float32, the default complex dtype is complex64. For float64, it is complex128. For bfloat16, an exception will be raised because there is no corresponding complex type for bfloat16. 2. Infer the dtype for tensors constructed using Python floats or complex Python numbers. See examples below. 3. Determine the result of type promotion between bool and integer tensors and Python floats and complex Python numbers. Args: d (:class:`tensorplay.dtype`): the floating point dtype to make the default. Example: >>> # xdoctest: +SKIP("Other tests may have changed the default type. Can we reset it?") >>> # initial default for floating point is float32 >>> # Python floats are interpreted as float32 >>> tensorplay.tensor([1.2, 3]).dtype tensorplay.float32 >>> # initial default for floating point is complex64 >>> # Complex Python numbers are interpreted as complex64 >>> tensorplay.tensor([1.2, 3j]).dtype tensorplay.complex64 >>> tensorplay.set_default_dtype(tensorplay.float64) >>> # Python floats are now interpreted as float64 >>> tensorplay.tensor([1.2, 3]).dtype # a new floating point tensor tensorplay.float64 >>> # Complex Python numbers are now interpreted as complex128 >>> tensorplay.tensor([1.2, 3j]).dtype # a new complex tensor tensorplay.complex128 >>> tensorplay.set_default_dtype(tensorplay.float16) >>> # Python floats are now interpreted as float16 >>> tensorplay.tensor([1.2, 3]).dtype # a new floating point tensor tensorplay.float16 >>> # Complex Python numbers are now interpreted as complex128 >>> tensorplay.tensor([1.2, 3j]).dtype # a new complex tensor tensorplay.complex32 """ _C._set_default_dtype(d) global _default_dtype _default_dtype = d [docs] def use_deterministic_algorithms( mode: builtins.bool, *, warn_only: builtins.bool = False, ) -> None: r"""Sets whether TensorPlay operations must use "deterministic" algorithms. That is, algorithms which, given the same input, and when run on the same software and hardware, always produce the same output. When enabled, operations will use deterministic algorithms when available, and if only nondeterministic algorithms are available they will throw a :class:`RuntimeError` when called. .. note:: This setting alone is not always enough to make an application reproducible. Refer to :ref:`reproducibility` for more information. .. note:: :func:`tensorplay.set_deterministic_debug_mode` offers an alternative interface for this feature. Note that deterministic operations tend to have worse performance than nondeterministic operations. .. note:: This flag does not detect or prevent nondeterministic behavior caused by calling an inplace operation on a tensor with an internal memory overlap or by giving such a tensor as the :attr:`out` argument for an operation. In these cases, multiple writes of different data may target a single memory location, and the order of writes is not guaranteed. Args: mode (:class:`bool`): If True, makes potentially nondeterministic operations switch to a deterministic algorithm or throw a runtime error. If False, allows nondeterministic operations. Keyword args: warn_only (:class:`bool`, optional): If True, operations that do not have a deterministic implementation will throw a warning instead of an error. Default: ``False`` Example:: >>> # xdoctest: +SKIP >>> tensorplay.use_deterministic_algorithms(True) """ # This flag has no compiler-specific counterpart. _C._set_deterministic_algorithms(mode, warn_only=warn_only) [docs] def are_deterministic_algorithms_enabled() -> builtins.bool: r"""Returns True if the global deterministic flag is turned on. Refer to :func:`tensorplay.use_deterministic_algorithms` documentation for more details. """ return _C._get_deterministic_algorithms() [docs] def is_deterministic_algorithms_warn_only_enabled() -> builtins.bool: r"""Returns True if the global deterministic flag is set to warn only. Refer to :func:`tensorplay.use_deterministic_algorithms` documentation for more details. """ return _C._get_deterministic_algorithms_warn_only() [docs] def set_deterministic_debug_mode(debug_mode: builtins.int | str) -> None: r"""Sets the debug mode for deterministic operations. .. note:: This is an alternative interface for :func:`tensorplay.use_deterministic_algorithms`. Refer to that function's documentation for details about affected operations. Args: debug_mode(str or int): If "default" or 0, don't error or warn on nondeterministic operations. If "warn" or 1, warn on nondeterministic operations. If "error" or 2, error on nondeterministic operations. """ # NOTE: builtins.int is used here because int in this scope resolves # to tensorplay.int if not isinstance(debug_mode, (builtins.int, str)): raise TypeError(f"debug_mode must be str or int, but got {type(debug_mode)}") if isinstance(debug_mode, str): if debug_mode == "default": debug_mode = 0 elif debug_mode == "warn": debug_mode = 1 elif debug_mode == "error": debug_mode = 2 else: raise RuntimeError( "invalid value of debug_mode, expected one of `default`, " f"`warn`, `error`, but got {debug_mode}" ) if debug_mode == 0: _C._set_deterministic_algorithms(False) elif debug_mode == 1: _C._set_deterministic_algorithms(True, warn_only=True) elif debug_mode == 2: _C._set_deterministic_algorithms(True) else: raise RuntimeError( f"invalid value of debug_mode, expected 0, 1, or 2, but got {debug_mode}" ) [docs] def get_deterministic_debug_mode() -> builtins.int: r"""Returns the current value of the debug mode for deterministic operations. Refer to :func:`tensorplay.set_deterministic_debug_mode` documentation for more details. """ if _C._get_deterministic_algorithms(): if _C._get_deterministic_algorithms_warn_only(): return 1 else: return 2 else: return 0 [docs] def get_float32_matmul_precision() -> str: r"""Returns the current value of float32 matrix multiplication precision. Refer to :func:`tensorplay.set_float32_matmul_precision` documentation for more details. """ return _C.get_float32_matmul_precision() [docs] def set_float32_matmul_precision(precision: str) -> None: r"""Sets the internal precision of float32 matrix multiplications. Running float32 matrix multiplications in lower precision may significantly increase performance, and in some programs the loss of precision has a negligible impact. Supports three settings: * "highest", float32 matrix multiplications use the float32 datatype (24 mantissa bits with 23 bits explicitly stored) for internal computations. * "high", float32 matrix multiplications either use the TensorFloat32 datatype (10 mantissa bits explicitly stored) or treat each float32 number as the sum of two bfloat16 numbers (approximately 16 mantissa bits with 14 bits explicitly stored), if the appropriate fast matrix multiplication algorithms are available. Otherwise float32 matrix multiplications are computed as if the precision is "highest". See below for more information on the bfloat16 approach. * "medium", float32 matrix multiplications use the bfloat16 datatype (8 mantissa bits with 7 bits explicitly stored) for internal computations, if a fast matrix multiplication algorithm using that datatype internally is available. Otherwise float32 matrix multiplications are computed as if the precision is "high". .. [Henry2019] http://arxiv.org/abs/1904.06376 .. note:: This does not change the output dtype of float32 matrix multiplications, it controls how the internal computation of the matrix multiplication is performed. .. note:: This does not change the precision of convolution operations. Other flags, like `tensorplay.backends.cudnn.allow_tf32`, may control the precision of convolution operations. .. note:: This flag currently only affects one native device type: CUDA. If "high" or "medium" are set then the TensorFloat32 datatype will be used when computing float32 matrix multiplications, equivalent to setting `tensorplay.backends.cuda.matmul.allow_tf32 = True`. When "highest" (the default) is set then the float32 datatype is used for internal computations, equivalent to setting `tensorplay.backends.cuda.matmul.allow_tf32 = False`. Args: precision(str): can be set to "highest" (default), "high", or "medium" (see above). """ if not isinstance(precision, str): raise TypeError("set_float32_matmul_precision expects a str, " f"but got {type(precision)}") _C._set_float32_matmul_precision(precision) __all__.extend([ "inference_mode", "get_default_dtype", "set_default_dtype", "get_default_device", "set_default_device", "use_deterministic_algorithms", "are_deterministic_algorithms_enabled", "is_deterministic_algorithms_warn_only_enabled", "set_deterministic_debug_mode", "get_deterministic_debug_mode", "get_float32_matmul_precision", "set_float32_matmul_precision", "get_printoptions", "printoptions", ]) newaxis: None = None __all__.extend(["e", "pi", "nan", "inf", "newaxis"]) from tensorplay._tensor import Tensor # The _tensor_classes set is initialized by the call to initialize_python_bindings. _tensor_classes: set[type[Tensor]] = set() import tensorplay __all__.extend( name for name in dir(tensorplay) if isinstance(getattr(tensorplay, name), tensorplay.dtype) ) # needs to be after the above c++ bindings so we can overwrite from Python side from tensorplay import functional as functional from tensorplay.functional import * # Underscore ops are skipped by the star-import above; expose them # explicitly so tensorplay._standard_gamma / _sample_dirichlet / # _safe_softmax resolve. from tensorplay.functional import ( _standard_gamma as _standard_gamma, _sample_dirichlet as _sample_dirichlet, _safe_softmax as _safe_softmax, ) # Keep package-level dtype aliases and composite signatures after importing # the generated function surface. chalf = DType.complex32 from tensorplay._shape_funcs import ( block_diag as block_diag, broadcast_tensors as broadcast_tensors, meshgrid as meshgrid, tensordot as tensordot, ) # Re-pin the einsum shim: functional.py's thin wrapper predates the sublist # quantile/nanquantile/histogram need the same treatment: the generated # functional.py wrappers forward a raw Python-number `q` (the _C binding # requires a Tensor) and drop `histogram`'s `range` keyword, so the # hand-written _composite_funcs versions (scalar-q -> input-dtype Tensor from ._einsum import einsum as einsum from ._composite_funcs import quantile as quantile from ._composite_funcs import nanquantile as nanquantile from ._composite_funcs import histogram as histogram from ._finfo import finfo, iinfo from ._tensor_str import ( get_printoptions as get_printoptions, printoptions as printoptions, set_printoptions as set_printoptions, ) # Python's ``import *`` intentionally omits underscore-prefixed names, but # only generated ``_foreach_*`` wrappers; their implementation is still the # native dispatcher/backend and this block does not introduce a Python # composite operator. ``_amp_*`` dispatcher hooks follow the same rule for _foreach_name in dir(functional): if (_foreach_name.startswith("_foreach_") or _foreach_name.startswith("_amp_") or _foreach_name.startswith("_fused_") or _foreach_name.startswith("_nested_")): globals()[_foreach_name] = getattr(functional, _foreach_name) if _foreach_name not in __all__: __all__.append(_foreach_name) del _foreach_name # ------------------------------------------------------------------------- # ``from tensorplay.functional import *`` above, which shadows earlier defs). # # reduction returning a named tuple with ``values``/``indices``, and the # reduction faces only, so the binary face and the named-tuple contract are # restored here at the package boundary. # ------------------------------------------------------------------------- import collections as _collections from ._return_types import ( # noqa: E402 aminmax_return_type as _aminmax_return_type, cummax_return_type as _cummax_return_type, cummin_return_type as _cummin_return_type, kthvalue_return_type as _kthvalue_return_type, max_return_type as _max_return_type, median_return_type as _median_return_type, min_return_type as _min_return_type, mode_return_type as _mode_return_type, nanmedian_return_type as _nanmedian_return_type, sort_return_type as _sort_return_type, topk_return_type as _topk_return_type, ) [docs] def max(input, *args, dim=None, keepdim=False): # ``args`` carries the elementwise partner of the two-argument form, or # the ``(dim, keepdim)`` pair when a recorded graph replays the native # call form. ``dim``/``keepdim`` stay keyword-only for the public API. # A bare integer in the second slot is the reduction dimension. if len(args) == 2: result = _C.max(input, args[0], args[1]) if not isinstance(result, tuple): return result return _max_return_type(*result) other = args[0] if args else None if other is not None and not isinstance(other, builtins.int): return _C.maximum(input, other) if other is not None: dim = other if dim is not None: result = _C.max(input, dim, keepdim) # A graph value stands for the whole multi-output call: consumers # split the outputs downstream, so the named tuple is only # constructed for eager results. if not isinstance(result, tuple): return result return _max_return_type(*result) return functional.max(input) [docs] def min(input, *args, dim=None, keepdim=False): # Same dual-form convention as max above. if len(args) == 2: result = _C.min(input, args[0], args[1]) if not isinstance(result, tuple): return result return _min_return_type(*result) other = args[0] if args else None if other is not None and not isinstance(other, builtins.int): return _C.minimum(input, other) if other is not None: dim = other if dim is not None: result = _C.min(input, dim, keepdim) if not isinstance(result, tuple): return result return _min_return_type(*result) return functional.min(input) def _wrap_multi_output(type_factory, positional_dims=1): """Rebuild a multi-output native call so eager results carry field names. Graph values stand for the whole multi-output call and are returned untouched, mirroring the max/min convention above. """ def wrapper(input, *args, **kwargs): result = getattr(functional, wrapper.__name__)(input, *args, **kwargs) if not isinstance(result, tuple): return result return type_factory(*result) wrapper.__name__ = type_factory.__name__.removesuffix("_return_type") wrapper.__qualname__ = wrapper.__name__ return wrapper cummax = _wrap_multi_output(_cummax_return_type) cummin = _wrap_multi_output(_cummin_return_type) kthvalue = _wrap_multi_output(_kthvalue_return_type) mode = _wrap_multi_output(_mode_return_type) sort = _wrap_multi_output(_sort_return_type) topk = _wrap_multi_output(_topk_return_type) nanmedian = _wrap_multi_output(_nanmedian_return_type) def aminmax(input, *args, **kwargs): result = functional.aminmax(input, *args, **kwargs) if not isinstance(result, tuple): return result return _aminmax_return_type(*result) def median(input, *args, **kwargs): # Without a dimension the reduction collapses to a single value; the # dimension form returns the (values, indices) pair. result = functional.median(input, *args, **kwargs) if not isinstance(result, tuple): return result return _median_return_type(*result) # --------------------------------------------------------------------------- # Sparse legacy matmuls: each is a thin routing over the sparse namespace. # --------------------------------------------------------------------------- def spmm(mat1, mat2): """Matrix product of a sparse COO/CSR matrix ``mat1`` with dense ``mat2``.""" from tensorplay import sparse as _sparse return _sparse.mm(mat1, mat2) def dsmm(mat1, mat2): """Matrix product of a dense matrix ``mat1`` with sparse ``mat2``. The sparse backend has no transposed product yet, so ``mat2`` is materialized for this legacy face; prefer the ``@`` operator or the sparse namespace for large operands. """ return mat1 @ mat2.to_dense() def hsmm(mat1, mat2): """Hybrid matrix product: dense ``mat1`` by sparse ``mat2``, dense result.""" return dsmm(mat1, mat2) def saddmm(input, mat1, mat2, *, beta=1, alpha=1, out=None): """``beta * input + alpha * (mat1 @ mat2)`` with sparse ``mat1`` and dense ``mat2``; ``input`` is a sparse matrix. A zero ``beta`` drops ``input`` entirely, so the sparse operand never goes through a dense copy on that path. """ from tensorplay import sparse as _sparse if beta == 0: result = _sparse.mm(mat1, mat2) if alpha != 1: result = result * alpha else: result = _sparse.addmm(input.to_dense(), mat1, mat2, beta=beta, alpha=alpha) if out is not None: out.resize_as_(result).copy_(result) return out return result # --------------------------------------------------------------------------- # Autocast state queries/setters that round out the public surface. # --------------------------------------------------------------------------- def is_autocast_cpu_enabled() -> builtins.bool: r"""Reports whether autocast is enabled on CPU.""" return _C.is_autocast_enabled("cpu") def set_autocast_cpu_enabled(enabled: builtins.bool) -> None: r""" Enables or disables autocast for the CPU. """ _C.set_autocast_enabled("cpu", enabled) def set_autocast_cpu_dtype(dtype) -> None: r"""Sets the autocast dtype for the CPU to ``dtype``.""" if dtype not in (bfloat16, float16): raise ValueError( "Supported CPU autocast dtypes are bfloat16 and float16; " f"got {dtype}" ) _C.set_autocast_dtype("cpu", dtype) def set_autocast_gpu_dtype(dtype) -> None: r"""Sets the autocast dtype for CUDA to ``dtype``.""" if dtype not in (bfloat16, float16): raise ValueError( "Supported CUDA autocast dtypes are bfloat16 and float16; " f"got {dtype}" ) _C.set_autocast_dtype("cuda", dtype) # --------------------------------------------------------------------------- # Futures and device-module helpers. # --------------------------------------------------------------------------- def wait(future): """Waits on the future and returns its value when it completes.""" return future.wait() @functools.cache def get_device_module(device=None): """Returns the module associated with a given device (e.g. ``tensorplay.device('cuda')``, ``"cpu"``). If no device is given, return the module for the current accelerator, or CPU if no accelerator is present. """ if isinstance(device, Device): device_module_name = device.type elif isinstance(device, str): device_module_name = Device(device).type elif device is None: from . import accelerator as _accelerator _acc = _accelerator.current_accelerator() device_module_name = _acc.type if _acc is not None else "cpu" else: raise RuntimeError( f"Invalid value of device '{device}', expect tensorplay.device, str, or None" ) device_module = globals().get(device_module_name, None) if device_module is None: try: device_module = importlib.import_module(f".{device_module_name}", __name__) except ModuleNotFoundError: device_module = None if device_module is None: raise RuntimeError( f"Device '{device_module_name}' does not have a corresponding module " f"registered as 'tensorplay.{device_module_name}'." ) return device_module [docs] def gradient(input, *, spacing=None, dim=None, edge_order=1): """ The native binding only takes ``Tensor[]`` spacing / ``int[]`` dim; materialize python numbers into scalar tensors typed like ``input`` (a single scalar applies to every dimension). """ dims = [dim] if isinstance(dim, builtins.int) else ( list(dim) if dim is not None else list(builtins.range(input.dim()))) if spacing is None: sp = [] elif isinstance(spacing, (builtins.int, builtins.float)): sp = [full([], spacing, dtype=input.dtype) for _ in dims] elif isinstance(spacing, Tensor): # A single coordinate tensor (numpy-style) applies to every # requested dimension; the kernel broadcasts numel()==n vs ==1. sp = [spacing] else: sp = [v if isinstance(v, Tensor) else full([], v, dtype=input.dtype) for v in spacing] return _C.gradient(self=input, spacing=sp, dim=dims, edge_order=edge_order) ################################################################################ # Import most common subpackages ################################################################################ # Use the redundant form so that type checkers know that these are a part of # the public API. The "regular" import lines are there solely for the runtime # side effect of adding to the imported module's members for other users. # needs to be before import tensorplay.nn as nn to avoid circular dependencies from tensorplay.autograd import ( enable_grad as enable_grad, inference_mode as inference_mode, no_grad as no_grad, set_grad_enabled as set_grad_enabled, is_grad_enabled as is_grad_enabled, ) from tensorplay import ( __config__ as __config__, __future__ as __future__, autograd as autograd, backends as backends, cuda as cuda, futures as futures, hub as hub, multiprocessing as multiprocessing, nn as nn, optim as optim, overrides as overrides, types as types, utils as utils, ) if TYPE_CHECKING: # Import the following modules during type checking to enable code intelligence features, # such as auto-completion for tools like pylance, even when these modules are not explicitly # imported in user code. from tensorplay import ( export as export, func as func, quantum as quantum, ) else: _lazy_modules = { "accelerator", "audio", "cpu", "export", "func", "distributions", "fft", "linalg", "ao", "contrib", "masked", "metrics", "monitor", "nested", "package", "quantum", "return_types", "signal", "snn", "sparse", "special", "storage", "testing", "vision", } # Public single names that live inside a lazily loaded module. The # structural control flow operators land here: their modules reach back # into the package surface through the graph and nn layers, and the # symbolic-shape layer they build on pulls in sympy, which is by far the # most expensive import in the package. Nothing outside this file # touches the operators during initialization, so first use pays for # them. _lazy_attrs = { "quantized_lstm": ("ao.nn.quantized._ops", "quantized_lstm"), "quantized_gru": ("ao.nn.quantized._ops", "quantized_gru"), "vmap": ("func", "vmap"), "cond": ("_higher_order_ops", "cond"), "while_loop": ("_higher_order_ops", "while_loop"), "OutOfMemoryError": ("cuda", "OutOfMemoryError"), } def __getattr__(name): # Lazy modules if name in _lazy_modules: return importlib.import_module(f".{name}", __name__) if name in _lazy_attrs: module, attr = _lazy_attrs[name] value = getattr(importlib.import_module(f".{module}", __name__), attr) # Bind into the package globals so the eager operators cost a # dict lookup per call rather than an import machinery round trip. globals()[name] = value return value raise AttributeError(f"module '{__name__}' has no attribute '{name}'") def _import_device_backends(): """ Leverage the Python plugin mechanism to load out-of-the-tree device extensions. """ from importlib.metadata import entry_points group_name = "tensorplay.backends" if sys.version_info < (3, 10): backend_extensions = entry_points().get(group_name, ()) else: backend_extensions = entry_points(group=group_name) for backend_extension in backend_extensions: try: # Load the extension entrypoint = backend_extension.load() # Call the entrypoint entrypoint() except Exception as err: raise RuntimeError( f"Failed to load the backend extension: {backend_extension.name}. " f"You can disable extension auto-loading with TENSORPLAY_DEVICE_BACKEND_AUTOLOAD=0." ) from err def _is_device_backend_autoload_enabled() -> builtins.bool: """ Whether autoloading out-of-the-tree device extensions is enabled. The switch depends on the value of the environment variable `TENSORPLAY_DEVICE_BACKEND_AUTOLOAD`. Returns: bool: Whether to enable autoloading the extensions. Enabled by default. Examples: >>> tensorplay._is_device_backend_autoload_enabled() True """ # enabled by default return os.getenv("TENSORPLAY_DEVICE_BACKEND_AUTOLOAD", "1") == "1" def _check(cond, msg=None): """ Assert a condition that may be backed by a symbolic expression. Under eager execution a plain check: a false condition raises ``RuntimeError`` carrying ``msg``. Symbolic backends may turn the check into a deferred assertion instead of an immediate failure. ``msg`` may be given as something to be called rather than as the message itself, so that a caller whose message costs something to write does not pay for it on the way to a check that holds. A message that is a callable is called only once the check has failed, and what it returns is the message. """ if not cond: if callable(msg): msg = msg() raise RuntimeError(msg if msg is not None else "Expected cond to be true") def _as_tensor_fullprec(t): """ Like tensorplay.as_tensor, but when given Python data types it will keep them in full precision. Used when a captured region converts Python values to tensors and the exact scalar widths must survive. """ ty = type(t) if ty is builtins.float: return tensorplay.as_tensor(t, dtype=tensorplay.float64) elif ty is builtins.int: return tensorplay.as_tensor(t, dtype=tensorplay.int64) else: return tensorplay.as_tensor(t) # --------------------------------------------------------------------------- # Public-surface presentation pass. Runs last so it sees every binding the # extension installed. # --------------------------------------------------------------------------- class _CleanFunction: """Python shim over a bound C entry whose own object is an extension helper type. Only the object identity differs: ``__call__`` forwards to the original, and the docstring is carried over.""" # The identity attributes live in the instance __dict__: ``__doc__`` and # ``__module__`` are class attributes and cannot be slots. __slots__ = ("_fn", "__dict__") def __init__(self, fn, name): self._fn = fn self.__name__ = self.__qualname__ = name self.__doc__ = fn.__doc__ self.__module__ = getattr(fn, "__module__", None) or "tensorplay" self.__wrapped__ = fn def __repr__(self): return f"" def __call__(self, *args, **kwargs): return self._fn(*args, **kwargs) def _polish_public_surface(): # The package's own ``types`` submodule shadows the standard module here. import types as _stdlib_types # 1) Classes re-exported here but registered on the extension quote the # extension's module path in type errors and reprs ("tensorplay._C.DType" # in `unsupported operand type(s)`); retitle them to the public package. for name, obj in list(globals().items()): if name.startswith("_") or not isinstance(obj, type): continue # Match on the exported name: a class may be registered on the # extension under an internal spell (``TensorBase``) while the # public alias differs, so its own __qualname__ is not a usable # lookup key. Rewriting __qualname__ to the exported spelling # keeps ``tensorplay.`` the one name everywhere (generated # code relies on that resolution). if ( obj.__module__ in ("tensorplay", "tensorplay._C") and globals().get(name) is obj ): # immutable extension types (e.g. Size) reject retitling; the # exported spelling is already theirs to keep then try: if obj.__module__ != "tensorplay": obj.__module__ = "tensorplay" if obj.__qualname__ != name: obj.__qualname__ = name except TypeError: pass # 2) Extension functions carry a binding-helper object as their bound # self; repr() and __qualname__ then quote that helper's C type name # ("PyCapsule"). Wrap those in a shim with a clean repr. Regular # builtins, extension classes and non-callables pass through # untouched. for module in list(sys.modules.values()): if module is None or not isinstance(module, _stdlib_types.ModuleType): continue module_name = getattr(module, "__name__", "") if not module_name.startswith("tensorplay") or "._" in module_name: continue for name, obj in list(vars(module).items()): # The shim only targets C functions, so submodules and other # non-callables are skipped. Probing them is not just wasted # work: attribute access on lazy proxies such as the cuFFT # plan-cache manager runs device init and raises on CPU-only # builds, and getattr's default only suppresses AttributeError. if ( name.startswith("_") or isinstance(obj, (type, _stdlib_types.ModuleType)) or not callable(obj) ): continue try: qualname = getattr(obj, "__qualname__", None) except Exception: # a callable proxy may still run arbitrary code on access continue if not isinstance(qualname, str) or not qualname.startswith("PyCapsule."): continue vars(module)[name] = _CleanFunction(obj, name) # `_import_device_backends` should be kept at the end to ensure # all the other functions in this module that may be accessed by # an autoloaded backend are defined if _is_device_backend_autoload_enabled(): _import_device_backends() _polish_public_surface() # Structural control flow operators; map and scan stay under # tensorplay._higher_order_ops. ``cond`` and ``while_loop`` bind through the # package ``__getattr__`` above rather than an import here: the operator # modules reach back into the package surface through the graph and nn layers, # so binding them during initialization both closes that cycle and forces the # symbolic-shape machinery (and sympy with it) into every session that never # runs a control flow operator. ```