Source code for treetensor.torch.funcs.construct

import torch
from treevalue import TreeValue
from treevalue.tree.common import TreeStorage

from .base import doc_from_base, func_treelize
from ..stream import stream_call
from ...utils import args_mapping

__all__ = [
    'tensor', 'as_tensor', 'clone',
    'zeros', 'zeros_like',
    'randn', 'randn_like',
    'randint', 'randint_like',
    'ones', 'ones_like',
    'full', 'full_like',
    'empty', 'empty_like',
]

args_treelize = args_mapping(lambda i, x: TreeValue(x) if isinstance(x, (dict, TreeStorage, TreeValue)) else x)


[docs]@doc_from_base() @args_treelize @func_treelize() def tensor(data, *args, **kwargs): """ In ``treetensor``, you can create a tree tensor with simple data structure. Examples:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.tensor(True) # the same as torch.tensor(True) tensor(True) >>> ttorch.tensor([1, 2, 3]) # the same as torch.tensor([1, 2, 3]) tensor([1, 2, 3]) >>> ttorch.tensor({'a': 1, 'b': [1, 2, 3], 'c': [[True, False], [False, True]]}) <Tensor 0x7ff363bbcc50> ├── a --> tensor(1) ├── b --> tensor([1, 2, 3]) └── c --> tensor([[ True, False], [False, True]]) """ return stream_call(torch.tensor, data, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def as_tensor(data, *args, **kwargs): """ Convert the data into a :class:`treetensor.torch.Tensor` or :class:`torch.Tensor`. Examples:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.as_tensor(True) tensor(True) >>> ttorch.as_tensor([1, 2, 3], dtype=torch.float32) tensor([1., 2., 3.]) >>> ttorch.as_tensor({ ... 'a': torch.tensor([1, 2, 3]), ... 'b': {'x': [[4, 5], [6, 7]]} ... }, dtype=torch.float32) <Tensor 0x7fc2b80c25c0> ├── a --> tensor([1., 2., 3.]) └── b --> <Tensor 0x7fc2b80c24e0> └── x --> tensor([[4., 5.], [6., 7.]]) """ return stream_call(torch.as_tensor, data, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @func_treelize() def clone(input, *args, **kwargs): """ In ``treetensor``, you can create a clone of the original tree with :func:`treetensor.torch.clone`. Examples:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.clone(torch.tensor([[1, 2], [3, 4]])) tensor([[1, 2], [3, 4]]) >>> ttorch.clone(ttorch.tensor({ ... 'a': [[1, 2], [3, 4]], ... 'b': {'x': [[5], [6], [7]]}, ... })) <Tensor 0x7f2a820ba5e0> ├── a --> tensor([[1, 2], │ [3, 4]]) └── b --> <Tensor 0x7f2a820aaf70> └── x --> tensor([[5], [6], [7]]) """ return stream_call(torch.clone, input, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def zeros(*args, **kwargs): """ In ``treetensor``, you can use ``zeros`` to create a tree of tensors with all zeros. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.zeros(2, 3) # the same as torch.zeros(2, 3) tensor([[0., 0., 0.], [0., 0., 0.]]) >>> ttorch.zeros({'a': (2, 3), 'b': {'x': (4, )}}) <Tensor 0x7f5f6ccf1ef0> ├── a --> tensor([[0., 0., 0.], │ [0., 0., 0.]]) └── b --> <Tensor 0x7f5fe0107208> └── x --> tensor([0., 0., 0., 0.]) """ return stream_call(torch.zeros, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @args_treelize @func_treelize() def zeros_like(input, *args, **kwargs): """ In ``treetensor``, you can use ``zeros_like`` to create a tree of tensors with all zeros like another tree. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.zeros_like(torch.randn(2, 3)) # the same as torch.zeros_like(torch.randn(2, 3)) tensor([[0., 0., 0.], [0., 0., 0.]]) >>> ttorch.zeros_like({ ... 'a': torch.randn(2, 3), ... 'b': {'x': torch.randn(4, )}, ... }) <Tensor 0x7ff363bb6128> ├── a --> tensor([[0., 0., 0.], │ [0., 0., 0.]]) └── b --> <Tensor 0x7ff363bb6080> └── x --> tensor([0., 0., 0., 0.]) """ return stream_call(torch.zeros_like, input, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def randn(*args, **kwargs): """ In ``treetensor``, you can use ``randn`` to create a tree of tensors with numbers obey standard normal distribution. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.randn(2, 3) # the same as torch.randn(2, 3) tensor([[-0.8534, -0.5754, -0.2507], [ 0.0826, -1.4110, 0.9748]]) >>> ttorch.randn({'a': (2, 3), 'b': {'x': (4, )}}) <Tensor 0x7ff363bb6518> ├── a --> tensor([[ 0.5398, 0.7529, -2.0339], │ [-0.5722, -1.1900, 0.7945]]) └── b --> <Tensor 0x7ff363bb6438> └── x --> tensor([-0.7181, 0.1670, -1.3587, -1.5129]) """ return stream_call(torch.randn, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @args_treelize @func_treelize() def randn_like(input, *args, **kwargs): """ In ``treetensor``, you can use ``randn_like`` to create a tree of tensors with numbers obey standard normal distribution like another tree. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.randn_like(torch.ones(2, 3)) # the same as torch.randn_like(torch.ones(2, 3)) tensor([[ 1.8436, 0.2601, 0.9687], [ 1.6430, -0.1765, -1.1732]]) >>> ttorch.randn_like({ ... 'a': torch.ones(2, 3), ... 'b': {'x': torch.ones(4, )}, ... }) <Tensor 0x7ff3d6f3cb38> ├── a --> tensor([[-0.1532, 1.3965, -1.2956], │ [-0.0750, 0.6475, 1.1421]]) └── b --> <Tensor 0x7ff3d6f420b8> └── x --> tensor([ 0.1730, 1.6085, 0.6487, -1.1022]) """ return stream_call(torch.randn_like, input, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def randint(*args, **kwargs): """ In ``treetensor``, you can use ``randint`` to create a tree of tensors with numbers in an integer range. Examples:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.randint(10, (2, 3)) # the same as torch.randint(10, (2, 3)) tensor([[3, 4, 5], [4, 5, 5]]) >>> ttorch.randint(10, {'a': (2, 3), 'b': {'x': (4, )}}) <Tensor 0x7ff363bb6438> ├── a --> tensor([[5, 3, 7], │ [8, 1, 8]]) └── b --> <Tensor 0x7ff363bb6240> └── x --> tensor([8, 8, 2, 4]) """ return stream_call(torch.randint, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @args_treelize @func_treelize() def randint_like(input, *args, **kwargs): """ In ``treetensor``, you can use ``randint_like`` to create a tree of tensors with numbers in an integer range. Examples:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.randint_like(torch.ones(2, 3), 10) # the same as torch.randint_like(torch.ones(2, 3), 10) tensor([[0., 5., 0.], [2., 0., 9.]]) >>> ttorch.randint_like({ ... 'a': torch.ones(2, 3), ... 'b': {'x': torch.ones(4, )}, ... }, 10) <Tensor 0x7ff363bb6748> ├── a --> tensor([[3., 6., 1.], │ [8., 9., 5.]]) └── b --> <Tensor 0x7ff363bb6898> └── x --> tensor([4., 4., 7., 1.]) """ return stream_call(torch.randint_like, input, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def ones(*args, **kwargs): """ In ``treetensor``, you can use ``ones`` to create a tree of tensors with all ones. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.ones(2, 3) # the same as torch.ones(2, 3) tensor([[1., 1., 1.], [1., 1., 1.]]) >>> ttorch.ones({'a': (2, 3), 'b': {'x': (4, )}}) <Tensor 0x7ff363bb6eb8> ├── a --> tensor([[1., 1., 1.], │ [1., 1., 1.]]) └── b --> <Tensor 0x7ff363bb6dd8> └── x --> tensor([1., 1., 1., 1.]) """ return stream_call(torch.ones, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @args_treelize @func_treelize() def ones_like(input, *args, **kwargs): """ In ``treetensor``, you can use ``ones_like`` to create a tree of tensors with all ones like another tree. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.ones_like(torch.randn(2, 3)) # the same as torch.ones_like(torch.randn(2, 3)) tensor([[1., 1., 1.], [1., 1., 1.]]) >>> ttorch.ones_like({ ... 'a': torch.randn(2, 3), ... 'b': {'x': torch.randn(4, )}, ... }) <Tensor 0x7ff363bbc320> ├── a --> tensor([[1., 1., 1.], │ [1., 1., 1.]]) └── b --> <Tensor 0x7ff363bbc240> └── x --> tensor([1., 1., 1., 1.]) """ return stream_call(torch.ones_like, input, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def full(*args, **kwargs): """ In ``treetensor``, you can use ``ones`` to create a tree of tensors with the same value. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.full((2, 3), 2.3) # the same as torch.full((2, 3), 2.3) tensor([[2.3000, 2.3000, 2.3000], [2.3000, 2.3000, 2.3000]]) >>> ttorch.full({'a': (2, 3), 'b': {'x': (4, )}}, 2.3) <Tensor 0x7ff363bbc7f0> ├── a --> tensor([[2.3000, 2.3000, 2.3000], │ [2.3000, 2.3000, 2.3000]]) └── b --> <Tensor 0x7ff363bbc8d0> └── x --> tensor([2.3000, 2.3000, 2.3000, 2.3000]) """ return stream_call(torch.full, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @args_treelize @func_treelize() def full_like(input, *args, **kwargs): """ In ``treetensor``, you can use ``ones_like`` to create a tree of tensors with all the same value of like another tree. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.full_like(torch.randn(2, 3), 2.3) # the same as torch.full_like(torch.randn(2, 3), 2.3) tensor([[2.3000, 2.3000, 2.3000], [2.3000, 2.3000, 2.3000]]) >>> ttorch.full_like({ ... 'a': torch.randn(2, 3), ... 'b': {'x': torch.randn(4, )}, ... }, 2.3) <Tensor 0x7ff363bb6cf8> ├── a --> tensor([[2.3000, 2.3000, 2.3000], │ [2.3000, 2.3000, 2.3000]]) └── b --> <Tensor 0x7ff363bb69e8> └── x --> tensor([2.3000, 2.3000, 2.3000, 2.3000]) """ return stream_call(torch.full_like, input, *args, **kwargs)
[docs]@doc_from_base() @args_treelize @func_treelize() def empty(*args, **kwargs): """ In ``treetensor``, you can use ``ones`` to create a tree of tensors with the uninitialized values. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.empty(2, 3) # the same as torch.empty(2, 3) tensor([[-1.3267e-36, 3.0802e-41, 2.3000e+00], [ 2.3000e+00, 2.3000e+00, 2.3000e+00]]) >>> ttorch.empty({'a': (2, 3), 'b': {'x': (4, )}}) <Tensor 0x7ff363bb6080> ├── a --> tensor([[-3.6515e+14, 4.5900e-41, -1.3253e-36], │ [ 3.0802e-41, 2.3000e+00, 2.3000e+00]]) └── b --> <Tensor 0x7ff363bb66d8> └── x --> tensor([-3.6515e+14, 4.5900e-41, -3.8091e-38, 3.0802e-41]) """ return stream_call(torch.empty, *args, **kwargs)
# noinspection PyShadowingBuiltins
[docs]@doc_from_base() @args_treelize @func_treelize() def empty_like(input, *args, **kwargs): """ In ``treetensor``, you can use ``ones_like`` to create a tree of tensors with all the uninitialized values of like another tree. Example:: >>> import torch >>> import treetensor.torch as ttorch >>> ttorch.empty_like(torch.randn(2, 3)) # the same as torch.empty_like(torch.randn(2, 3), 2.3) tensor([[-3.6515e+14, 4.5900e-41, -1.3266e-36], [ 3.0802e-41, 4.4842e-44, 0.0000e+00]]) >>> ttorch.empty_like({ ... 'a': torch.randn(2, 3), ... 'b': {'x': torch.randn(4, )}, ... }) <Tensor 0x7ff363bbc780> ├── a --> tensor([[-3.6515e+14, 4.5900e-41, -3.6515e+14], │ [ 4.5900e-41, 1.1592e-41, 0.0000e+00]]) └── b --> <Tensor 0x7ff3d6f3cb38> └── x --> tensor([-1.3267e-36, 3.0802e-41, -3.8049e-38, 3.0802e-41]) """ return stream_call(torch.empty_like, input, *args, **kwargs)