Caffe2 - Python API A deep learning, cross platform ML framework
__init__.py
1 # The Tensor classes are added to this module by python_tensor.cpp
2 import torch
3
4 __all__ = [
6  'mm',
7  'sum',
8 ]
9
10
11 def addmm(mat, mat1, mat2, beta=1, alpha=1):
12  r"""
13  This function does exact same thing as :func:torch.addmm in the forward,
14  except that it supports backward for sparse matrix :attr:mat1. :attr:mat1
15  need to have sparse_dim = 2. Note that the gradients of :attr:mat1 is a
16  coalesced sparse tensor.
17
18  Args:
19  mat (Tensor): a dense matrix to be added
20  mat1 (SparseTensor): a sparse matrix to be multiplied
21  mat2 (Tensor): a dense matrix be multiplied
22  beta (Number, optional): multiplier for :attr:mat (:math:\beta)
23  alpha (Number, optional): multiplier for :math:mat1 @ mat2 (:math:\alpha)
24  """
26
27
28 def mm(mat1, mat2):
29  r"""
30  Performs a matrix multiplication of the sparse matrix :attr:mat1
31  and dense matrix :attr:mat2. Similar to :func:torch.mm, If :attr:mat1 is a
32  :math:(n \times m) tensor, :attr:mat2 is a :math:(m \times p) tensor, out will be a
33  :math:(n \times p) dense tensor. :attr:mat1 need to have sparse_dim = 2.
34  This function also supports backward for both matrices. Note that the gradients of
35  :attr:mat1 is a coalesced sparse tensor.
36
37  Args:
38  mat1 (SparseTensor): the first sparse matrix to be multiplied
39  mat2 (Tensor): the second dense matrix to be multiplied
40
41  Example::
42
43  >>> a = torch.randn(2, 3).to_sparse().requires_grad_(True)
44  >>> a
45  tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
46  [0, 1, 2, 0, 1, 2]]),
47  values=tensor([ 1.5901, 0.0183, -0.6146, 1.8061, -0.0112, 0.6302]),
48  size=(2, 3), nnz=6, layout=torch.sparse_coo, requires_grad=True)
49
50  >>> b = torch.randn(3, 2, requires_grad=True)
51  >>> b
52  tensor([[-0.6479, 0.7874],
53  [-1.2056, 0.5641],
55
56  >>> y = torch.sparse.mm(a, b)
57  >>> y
58  tensor([[-0.3323, 1.8723],
60  >>> y.sum().backward()
62  tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
63  [0, 1, 2, 0, 1, 2]]),
64  values=tensor([ 0.1394, -0.6415, -2.1639, 0.1394, -0.6415, -2.1639]),
65  size=(2, 3), nnz=6, layout=torch.sparse_coo)
66  """
68
69
70 def sum(input, dim=None, dtype=None):
71  r"""
72  Returns the sum of each row of SparseTensor :attr:input in the given
73  dimensions :attr:dim. If :attr:dim is a list of dimensions,
74  reduce over all of them. When sum over all sparse_dim, this method
75  returns a Tensor instead of SparseTensor.
76
77  All summed :attr:dim are squeezed (see :func:torch.squeeze), resulting an output
78  tensor having :attr:dim fewer dimensions than :attr:input.
79
80  During backward, only gradients at nnz locations of :attr:input
81  will propagate back. Note that the gradients of :attr:input is coalesced.
82
83  Args:
84  input (Tensor): the input SparseTensor
85  dim (int or tuple of ints): a dimension or a list of dimensions to reduce. Default: reduce
86  over all dims.
87  dtype (:class:torch.dtype, optional): the desired data type of returned Tensor.
88  Default: dtype of :attr:input.
89
90  Example::
91
92  >>> nnz = 3
93  >>> dims = [5, 5, 2, 3]
94  >>> I = torch.cat([torch.randint(0, dims[0], size=(nnz,)),
95  torch.randint(0, dims[1], size=(nnz,))], 0).reshape(2, nnz)
96  >>> V = torch.randn(nnz, dims[2], dims[3])
97  >>> size = torch.Size(dims)
98  >>> S = torch.sparse_coo_tensor(I, V, size)
99  >>> S
100  tensor(indices=tensor([[2, 0, 3],
101  [2, 4, 1]]),
102  values=tensor([[[-0.6438, -1.6467, 1.4004],
103  [ 0.3411, 0.0918, -0.2312]],
104
105  [[ 0.5348, 0.0634, -2.0494],
106  [-0.7125, -1.0646, 2.1844]],
107
108  [[ 0.1276, 0.1874, -0.6334],
109  [-1.9682, -0.5340, 0.7483]]]),
110  size=(5, 5, 2, 3), nnz=3, layout=torch.sparse_coo)
111
112  # when sum over only part of sparse_dims, return a SparseTensor
113  >>> torch.sparse.sum(S, [1, 3])
114  tensor(indices=tensor([[0, 2, 3]]),
115  values=tensor([[-1.4512, 0.4073],
116  [-0.8901, 0.2017],
117  [-0.3183, -1.7539]]),
118  size=(5, 2), nnz=3, layout=torch.sparse_coo)
119
120  # when sum over all sparse dim, return a dense Tensor
121  # with summed dims squeezed
122  >>> torch.sparse.sum(S, [0, 1, 3])
123  tensor([-2.6596, -1.1450])
124  """
125  if dtype is None:
126  if dim is not None: