def new_leaf(size, kwargs): """Return a leaf tensor from optional kwargs.""" device = kwargs.get('device', cpp.device()) return Tensor(*size, dtype=kwargs.get('dtype', 'float32'), device=cpp.device() if device is None else device, requires_grad=kwargs.get('requires_grad', False))
def empty(*size, dtype=None, device=None, requires_grad=False): """Return a tensor filled with uninitialized data. Parameters ---------- size : int... The sizes of output tensor. dtype : str, optional The optional data type. device : dragon.vm.torch.device, optional The optional device option. requires_grad : bool, optional, default=False Whether to compute the gradient if necessary. Returns ------- dragon.vm.torch.Tensor The output tensor. """ return Tensor( *size, dtype=dtype if dtype else 'float32', device=cpp.device() if device is None else device, requires_grad=requires_grad, )
def zeros(*size, out=None, dtype='float32', device=None, requires_grad=False): r"""Return a tensor filled with zeros. .. math:: \text{out} \leftarrow 0 Parameters ---------- size : int... The output tensor shape. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='float32' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ size = nest.flatten(size) device = out.device if out else (device or cpp.device()) out = Function.apply('Fill', device, [], outputs=[out], dtype=dtype, value=0.0, ndim=len(size), dims=size) out._requires_grad = requires_grad return out
def tensor(data, dtype=None, device=None, requires_grad=False): """Create a tensor initializing from the given data. Parameters ---------- data : array_like The data to initialize from. dtype : str, optional The optional data type. device : dragon.vm.torch.device, optional The optional device of returned tensor. requires_grad : bool, optional, default=False ``True`` to record gradient for returned tensor. Returns ------- dragon.vm.torch.Tensor The output tensor. """ array_data = numpy.array(data, copy=True) if dtype is None: dtype = str(array_data.dtype) else: array_data = array_data.astype(dtype) return Tensor( array_data, dtype=dtype, device=cpp.device() if device is None else device, requires_grad=requires_grad, )
def randn(*size, out=None, dtype='float32', device=None, requires_grad=False): """Return a tensor from the normal distribution of N(0, 1). Parameters ---------- size : int... The output tensor shape. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='float32' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ size = nest.flatten(size) device = out.device if out else (device or cpp.device()) out = Function.apply( 'RandomNormal', device, [], outputs=[out], dtype=dtype, mean=0.0, std=1.0, ndim=len(size), dims=size) out._requires_grad = requires_grad return out
def normal(mean, std, size, out=None): r"""Return a tensor initialized from the normal distribution. .. math:: \text{out} \sim \mathcal{N}(\mu, \sigma^{2}) Parameters ---------- mean : number The value to :math:`\mu`. std : number The value to :math:`\sigma`. size : Sequence[int] The output tensor shape. out : dragon.vm.torch.Tensor, optional The output tensor. Returns ------- dragon.vm.torch.Tensor The output tensor. """ dtype = out.dtype if out else 'float32' device = out.device if out else cpp.device() return Function.apply( 'RandomNormal', device, [], outputs=[out], dtype=dtype, mean=float(mean), std=float(std), ndim=len(size), dims=size)
def uniform(low, high, size, out=None): r"""Return a tensor initialized from the uniform distribution. .. math:: \text{out} \sim \mathcal{U}(\alpha, \beta) Parameters ---------- low : number The value to :math:`\alpha`. high : number The value to :math:`\beta`. size : Sequence[int] The output tensor shape. out : dragon.vm.torch.Tensor, optional The output tensor. Returns ------- dragon.vm.torch.Tensor The output tensor. """ dtype = out.dtype if out else 'float32' device = out.device if out else cpp.device() return Function.apply( 'RandomUniform', device, [], outputs=[out], dtype=dtype, low=float(low), high=float(high), ndim=len(size), dims=size)
def randperm(n, out=None, dtype='int64', device=None, requires_grad=False): """Return a tensor with value in the permuted range. Specify ``n`` to determine an interval :math:`[0, n)`: ```python print(torch.randperm(4)) ``` Parameters ---------- n: number The end of interval. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='int64' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ device = out.device if out else (device or cpp.device()) out = Function.apply( 'Permutation', device, [], outputs=[out], dtype=dtype, limit=n) out._requires_grad = requires_grad return out
def unify_devices(tensors, key='Inputs'): """Return a device unified from tensors.""" types, indices = [t._device.type for t in tensors], [0] if len(set(types)) != 1: raise ValueError('{} from different device type: [{}].'.format( key, ', '.join(types))) if types[0] == 'cuda': indices = [t._device.index for t in tensors] if len(set(indices)) != 1: raise ValueError('{} from different cuda device: [{}].'.format( key, ', '.join([str(d) for d in indices]))) return cpp.device(types[0], indices[0])
def eye( n, m=None, out=None, dtype='float32', device=None, requires_grad=False, ): r"""Return a tensor constructed as the identity matrix. .. math:: \text{out} \leftarrow \text{diag}(1, 1, ..., 1) The rows and cols of matrix are determined by ``n`` and ``m``: ```python print(torch.eye(2)) # [[1., 0.], [0., 1.]] print(torch.eye(2, 3)) # [[1., 0., 0.], [0., 1., 0.]] ``` Parameters ---------- n : int The number output rows. m : int, optional The number output cols. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='float32' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ m = n if m is None else m device = out.device if out else (device or cpp.device()) out = Function.apply('Eye', device, [], outputs=[out], dtype=dtype, ndim=2, dims=(n, m)) out._requires_grad = requires_grad return out
def full( size, fill_value, out=None, dtype='int64', device=None, requires_grad=False, ): """Return a tensor filled with a scalar. Examples: ```python print(torch.full((1, 2), 1)) # [[1, 1]] ``` Parameters ---------- size : int... The output shape. fill_value : number The scalar to fill. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='int64' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ size = nest.flatten(size) device = out.device if out else (device or cpp.device()) out = Function.apply('Fill', device, [], outputs=[out], dtype=dtype, value=float(fill_value), ndim=len(size), dims=size) out._requires_grad = requires_grad return out
def from_dlpack(dlpack): """Create a tensor sharing the dlpack data. Parameters ---------- dlpack : PyCapsule The capsule object of a dlpack tensor. Returns ------- dragon.vm.torch.Tensor The tensor with the dlpack data. """ default_ws = workspace.get_workspace() impl = default_ws.create_tensor(scope='DLPack').FromDLPack(dlpack) return Tensor(device=cpp.device(*impl.device), impl=impl, deleter=default_ws._handle_pool)
def full_like( input, fill_value, out=None, dtype='int64', device=None, requires_grad=False, ): """Return a tensor filled with a scalar with size as input. Examples: ```python print(torch.full_like(torch.zeros(1, 2), 1)) # [[1, 1]] ``` Parameters ---------- input : dragon.vm.torch.Tensor The tensor for indicating shape. fill_value : number The scalar to fill. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='int64' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ device = out.device if out else (device or cpp.device()) out = Function.apply('Fill', device, [input], dtype=dtype, value=float(fill_value)) out._requires_grad = requires_grad return out
def from_dlpack(dlpack): """Create a tensor sharing the dlpack data. Parameters ---------- dlpack : PyCapsule The capsule object of a dlpack tensor. Returns ------- dragon.vm.torch.Tensor The tensor with the dlpack data. """ current_ws = workspace.get_workspace() tensor = Tensor(device=None) tensor._gc = current_ws.collectors.TENSOR tensor._impl = current_ws.create_tensor( tensor._gc.alloc('${DLPACK}')).FromDLPack(dlpack) tensor._device = cpp.device(*tensor._impl.device) return tensor
def linspace( start, end, steps=100, out=None, dtype='int64', dim=0, device=None, requires_grad=False, ): r"""Generate evenly spaced values within intervals along the given dimension. Interval :math:`[\text{start}, \text{end})` is determined for ``steps`` values: ```python x = torch.linspace(2, 4, steps=3) # [2, 3, 4] ``` More than one intervals are accepted to generate N-d coordinates: ```python x = torch.linspace([1, 2], [3, 4], steps=3, dim=0) # [[1, 2], [2, 3], [3, 4]] y = torch.linspace([1, 2], [3, 4], steps=3, dim=1) # [[1, 2, 3], [2, 3, 4]] ``` Parameters ---------- start : Union[number, Sequence[number]] The start(s) of interval. end: Union[number, Sequence[number]] The ends(s) of interval. steps : int, optional, default=100 The number of values to generate. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='int64' The data type of output tensor. dim : int, optional, default=0 The dimension to generate values. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor The output tensor. """ starts = nest.flatten(start) ends = nest.flatten(end) size = [] if len(starts) > 1 or starts == start: size = [len(starts)] dim = dim if dim >= 0 else dim + len(size) + 1 size.insert(dim, steps) device = out.device if out else (device or cpp.device()) out = Function.apply('LinSpace', device, [], outputs=[out], dtype=dtype, axis=dim, ndim=len(size), num_intervals=len(starts), dims=size, start=starts, stop=ends) out._requires_grad = requires_grad return out
def arange( start, end=None, step=1, out=None, dtype='int64', device=None, requires_grad=False, ): """Return a tensor of evenly spaced values within a interval. Specify ``start`` and ``end`` to determine an interval: ```python print(torch.arange(2, 4)) # [2, 3] ``` If ``stop`` is **None**, interval :math:`[0, start)` will be taken instead: ```python print(torch.arange(5)) # [0, 1, 2, 3, 4] ``` Set ``delta`` to make the strides: ```python print(torch.arange(5, step=2)) # [0, 2, 4] ``` Parameters ---------- start : number The start of interval. end : number, optional, default=0 The stop of interval. step : number, optional, default=1 The spacing between two elements. out : dragon.vm.torch.Tensor, optional The output tensor. dtype : str, optional, default='int64' The data type of output tensor. device : dragon.vm.torch.device, optional The device of output tensor. requires_grad : bool, optional, default=False Record gradient for output tensor or not. Returns ------- dragon.vm.torch.Tensor A vector with evenly spaced elements. """ if end is None: slice_args = start, step else: slice_args = start, end, step device = out.device if out else (device or cpp.device()) out = Function.apply('Range', device, [], outputs=[out], dtype=dtype, num_args=len(slice_args), slice=slice_args) out._requires_grad = requires_grad return out