We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Negative dimension behaviour in torch.stack
Support for using a negative dimension in `torch.stack` currently works as follows:

```python
if dim < 0:
    dim += sequence[0].dim()
```

This is not consistent with the behaviour of `torch.unsqueeze`, which uses `dim += input.dim() + 1`. The unsqueeze behaviour is better since it is possible to use negative indexing to insert a new dimension as the last dimension.

## Example

`torch.stack` with `dim=-1` adds a new dimension in the second-to-last position. This is confusing. It is not possible to use negative indexing to add a dimension in the last position (ie can't achieve a 3x2 result).

```python
>>> a = torch.Tensor([1, 2, 3])
>>> b = torch.Tensor([4, 5, 6])
>>> torch.stack([a, b], -1)

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]
```

Contrast this to using `torch.unsqueeze` and `torch.cat` with `dim=-1`, which adds a new dimension in the last position:

```python
>>> a = torch.Tensor([1, 2, 3])
>>> b = torch.Tensor([4, 5, 6])
>>> torch.cat([a.unsqueeze(-1), b.unsqueeze(-1)], -1)

 1  4
 2  5
 3  6
[torch.FloatTensor of size 3x2]
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch/functional.py`
Content:
```
1 import torch
2 from ._utils import _range
3 from operator import mul
4 from functools import reduce
5 
6 __all__ = [
7     'split', 'chunk', 'stack', 'unbind', 'btriunpack', 'matmul',
8 ]
9 
10 
11 def split(tensor, split_size, dim=0):
12     """Splits the tensor into equally sized chunks (if possible).
13 
14     Last chunk will be smaller if the tensor size along a given dimension
15     is not divisible by ``split_size``.
16 
17     Arguments:
18         tensor (Tensor): tensor to split.
19         split_size (int): size of a single chunk.
20         dim (int): dimension along which to split the tensor.
21     """
22     if dim < 0:
23         dim += tensor.dim()
24     dim_size = tensor.size(dim)
25     num_splits = (dim_size + split_size - 1) // split_size
26     last_split_size = split_size - (split_size * num_splits - dim_size)
27 
28     def get_split_size(i):
29         return split_size if i < num_splits - 1 else last_split_size
30     return tuple(tensor.narrow(int(dim), int(i * split_size), int(get_split_size(i))) for i
31                  in _range(0, num_splits))
32 
33 
34 def chunk(tensor, chunks, dim=0):
35     """Splits a tensor into a number of chunks along a given dimension.
36 
37     Arguments:
38         tensor (Tensor): tensor to split.
39         chunks (int): number of chunks to return.
40         dim (int): dimension along which to split the tensor.
41     """
42     if dim < 0:
43         dim += tensor.dim()
44     split_size = (tensor.size(dim) + chunks - 1) // chunks
45     return split(tensor, split_size, dim)
46 
47 
48 def stack(sequence, dim=0, out=None):
49     """Concatenates sequence of tensors along a new dimension.
50 
51     All tensors need to be of the same size.
52 
53     Arguments:
54         sequence (Sequence): sequence of tensors to concatenate.
55         dim (int): dimension to insert. Has to be between 0 and the number
56             of dimensions of concatenated tensors (inclusive).
57     """
58     if len(sequence) == 0:
59         raise ValueError("stack expects a non-empty sequence of tensors")
60     if dim < 0:
61         dim += sequence[0].dim()
62     inputs = [t.unsqueeze(dim) for t in sequence]
63     if out is None:
64         return torch.cat(inputs, dim)
65     else:
66         return torch.cat(inputs, dim, out=out)
67 
68 
69 def unbind(tensor, dim=0):
70     """Removes a tensor dimension.
71 
72     Returns a tuple of all slices along a given dimension, already without it.
73 
74     Arguments:
75         tensor (Tensor): tensor to unbind.
76         dim (int): dimension to remove.
77     """
78     return tuple(tensor.select(dim, i) for i in _range(tensor.size(dim)))
79 
80 
81 def btriunpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True):
82     """Unpacks the data and pivots from a batched LU factorization (btrifact) of a tensor.
83 
84     Returns a tuple indexed by:
85       0: The pivots.
86       1: The L tensor.
87       2: The U tensor.
88 
89     Arguments:
90         LU_data (Tensor): The packed LU factorization data.
91         LU_pivots (Tensor): The packed LU factorization pivots.
92         unpack_data (bool): Flag indicating if the data should be unpacked.
93         unpack_pivots (bool): Flag indicating if the pivots should be unpacked.
94     """
95 
96     nBatch, sz, _ = LU_data.size()
97 
98     if unpack_data:
99         I_U = torch.triu(torch.ones(sz, sz)).type_as(LU_data).byte().unsqueeze(0).expand(nBatch, sz, sz)
100         I_L = 1 - I_U
101         L = LU_data.new(LU_data.size()).zero_()
102         U = LU_data.new(LU_data.size()).zero_()
103         I_diag = torch.eye(sz).type_as(LU_data).byte().unsqueeze(0).expand(nBatch, sz, sz)
104         L[I_diag] = 1.0
105         L[I_L] = LU_data[I_L]
106         U[I_U] = LU_data[I_U]
107     else:
108         L = U = None
109 
110     if unpack_pivots:
111         P = torch.eye(sz).type_as(LU_data).unsqueeze(0).repeat(nBatch, 1, 1)
112         for i in range(nBatch):
113             for j in range(sz):
114                 k = LU_pivots[i, j] - 1
115                 t = P[i, :, j].clone()
116                 P[i, :, j] = P[i, :, k]
117                 P[i, :, k] = t
118     else:
119         P = None
120 
121     return P, L, U
122 
123 
124 def matmul(tensor1, tensor2, out=None):
125     """Matrix product of two tensors.
126 
127     The behavior depends on the dimensionality of the tensors as follows:
128 
129     - If both tensors are 1-dimensional, the dot product (scalar) is returned.
130     - If both arguments are 2-dimensional, the matrix-matrix product is returned.
131     - If the first argument is 1-dimensional and the second argument is 2-dimensional,
132       a 1 is prepended to its dimension for the purpose of the matrix multiply.
133       After the matrix multiply, the prepended dimension is removed.
134     - If the first argument is 2-dimensional and the second argument is 1-dimensional,
135       the matrix-vector product is returned.
136     - If both arguments are at least 1-dimensional and at least one argument is
137       N-dimensional (where N > 2), then a batched matrix multiply is returned.  If the first
138       argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the
139       batched matrix multiply and removed after.  If the second argument is 1-dimensional, a
140       1 is appended to its dimension for the purpose of the batched matrix multiple and removed after.
141       The non-matrix (i.e. batch) dimensions are :ref:`broadcasted <broadcasting-semantics>` (and thus
142       must be broadcastable).  For example, if :attr:`tensor1` is a `j x 1 x n x m` Tensor
143       and :attr:`tensor2` is a `k x m x p` Tensor, :attr:`out` will be an `j x k x n x p` Tensor.
144 
145     .. note::
146 
147         The 1-dimensional dot product version of this function does not support an :attr:`out` parameter.
148 
149     Arguments:
150         tensor1 (Tensor): First tensor to be multiplied
151         tensor2 (Tensor): Second tensor to be multiplied
152         out (Tensor, optional): Output tensor
153     """
154     dim_tensor1 = tensor1.dim()
155     dim_tensor2 = tensor2.dim()
156     if dim_tensor1 == 1 and dim_tensor2 == 1:
157         if out is None:
158             return torch.dot(tensor1, tensor2)
159         else:
160             raise ValueError("out must be None for 1-d tensor matmul, returns a scalar")
161     if dim_tensor1 == 2 and dim_tensor2 == 1:
162         if out is None:
163             return torch.mv(tensor1, tensor2)
164         else:
165             return torch.mv(tensor1, tensor2, out=out)
166     elif dim_tensor1 == 1 and dim_tensor2 == 2:
167         if out is None:
168             return torch.mm(tensor1.unsqueeze(0), tensor2).squeeze_(0)
169         else:
170             return torch.mm(tensor1.unsqueeze(0), tensor2, out=out).squeeze_(0)
171     elif dim_tensor1 == 2 and dim_tensor2 == 2:
172         if out is None:
173             return torch.mm(tensor1, tensor2)
174         else:
175             return torch.mm(tensor1, tensor2, out=out)
176     elif dim_tensor1 >= 3 and (dim_tensor2 == 1 or dim_tensor2 == 2):
177         # optimization: use mm instead of bmm by folding tensor1's batch into
178         # its leading matrix dimension.
179 
180         if dim_tensor2 == 1:
181             tensor2 = tensor2.unsqueeze(-1)
182 
183         size1 = tensor1.size()
184         size2 = tensor2.size()
185         output_size = size1[:-1] + size2[-1:]
186 
187         # fold the batch into the first dimension
188         tensor1 = tensor1.contiguous().view(-1, size1[-1])
189 
190         if out is None or not out.is_contiguous():
191             output = torch.mm(tensor1, tensor2)
192         else:
193             output = torch.mm(tensor1, tensor2, out=out)
194 
195         output = output.view(output_size)
196 
197         if dim_tensor2 == 1:
198             output = output.squeeze(-1)
199 
200         if out is not None:
201             out.set_(output)
202             return out
203 
204         return output
205     elif (dim_tensor1 >= 1 and dim_tensor2 >= 1) and (dim_tensor1 >= 3 or dim_tensor2 >= 3):
206         # ensure each tensor size is at least 3-dimensional
207         tensor1_exp_size = torch.Size((1,) * max(3 - tensor1.dim(), 0) + tensor1.size())
208         # rhs needs to be a separate case since we can't freely expand 1s on the rhs, but can on lhs
209         if dim_tensor2 == 1:
210             tensor2 = tensor2.unsqueeze(1)
211         tensor2_exp_size = torch.Size((1,) * max(3 - tensor2.dim(), 0) + tensor2.size())
212 
213         # expand the batch portion (i.e. cut off matrix dimensions and expand rest)
214         expand_batch_portion = torch._C._infer_size(tensor1_exp_size[:-2], tensor2_exp_size[:-2])
215 
216         # flatten expanded batches
217         tensor1_expanded = tensor1.expand(*(expand_batch_portion + tensor1_exp_size[-2:])) \
218             .contiguous().view(reduce(mul, expand_batch_portion), *tensor1_exp_size[-2:])
219         tensor2_expanded = tensor2.expand(*(expand_batch_portion + tensor2_exp_size[-2:])) \
220             .contiguous().view(reduce(mul, expand_batch_portion), *tensor2_exp_size[-2:])
221 
222         # reshape batches back into result
223         total_expansion = expand_batch_portion + (tensor1_exp_size[-2], tensor2_exp_size[-1])
224 
225         def maybeSqueeze(tensor):
226             if dim_tensor1 == 1:
227                 return tensor.squeeze(-2)
228             elif dim_tensor2 == 1:
229                 return tensor.squeeze(-1)
230             else:
231                 return tensor
232 
233         if out is None or not out.is_contiguous():
234             output = torch.bmm(tensor1_expanded, tensor2_expanded)
235         else:
236             output = torch.bmm(tensor1_expanded, tensor2_expanded, out=out)
237 
238         output = maybeSqueeze(output.view(total_expansion))
239 
240         if out is not None:
241             out.set_(output)
242             return out
243 
244         return output
245 
246     raise ValueError("both arguments to __matmul__ need to be at least 1D, "
247                      "but they are {}D and {}D".format(dim_tensor1, dim_tensor2))
248
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch/functional.py b/torch/functional.py
--- a/torch/functional.py
+++ b/torch/functional.py
@@ -58,7 +58,7 @@
     if len(sequence) == 0:
         raise ValueError("stack expects a non-empty sequence of tensors")
     if dim < 0:
-        dim += sequence[0].dim()
+        dim += sequence[0].dim() + 1
     inputs = [t.unsqueeze(dim) for t in sequence]
     if out is None:
         return torch.cat(inputs, dim)

{"golden_diff": "diff --git a/torch/functional.py b/torch/functional.py\n--- a/torch/functional.py\n+++ b/torch/functional.py\n@@ -58,7 +58,7 @@\n     if len(sequence) == 0:\n         raise ValueError(\"stack expects a non-empty sequence of tensors\")\n     if dim < 0:\n-        dim += sequence[0].dim()\n+        dim += sequence[0].dim() + 1\n     inputs = [t.unsqueeze(dim) for t in sequence]\n     if out is None:\n         return torch.cat(inputs, dim)\n", "issue": "Negative dimension behaviour in torch.stack\nSupport for using a negative dimension in `torch.stack` currently works as follows:\r\n\r\n```python\r\nif dim < 0:\r\n    dim += sequence[0].dim()\r\n```\r\n\r\nThis is not consistent with the behaviour of `torch.unsqueeze`, which uses `dim += input.dim() + 1`. The unsqueeze behaviour is better since it is possible to use negative indexing to insert a new dimension as the last dimension.\r\n\r\n## Example\r\n\r\n`torch.stack` with `dim=-1` adds a new dimension in the second-to-last position. This is confusing. It is not possible to use negative indexing to add a dimension in the last position (ie can't achieve a 3x2 result).\r\n\r\n```python\r\n>>> a = torch.Tensor([1, 2, 3])\r\n>>> b = torch.Tensor([4, 5, 6])\r\n>>> torch.stack([a, b], -1)\r\n\r\n 1  2  3\r\n 4  5  6\r\n[torch.FloatTensor of size 2x3]\r\n```\r\n\r\nContrast this to using `torch.unsqueeze` and `torch.cat` with `dim=-1`, which adds a new dimension in the last position:\r\n\r\n```python\r\n>>> a = torch.Tensor([1, 2, 3])\r\n>>> b = torch.Tensor([4, 5, 6])\r\n>>> torch.cat([a.unsqueeze(-1), b.unsqueeze(-1)], -1)\r\n\r\n 1  4\r\n 2  5\r\n 3  6\r\n[torch.FloatTensor of size 3x2]\r\n```\n", "before_files": [{"content": "import torch\nfrom ._utils import _range\nfrom operator import mul\nfrom functools import reduce\n\n__all__ = [\n    'split', 'chunk', 'stack', 'unbind', 'btriunpack', 'matmul',\n]\n\n\ndef split(tensor, split_size, dim=0):\n    \"\"\"Splits the tensor into equally sized chunks (if possible).\n\n    Last chunk will be smaller if the tensor size along a given dimension\n    is not divisible by ``split_size``.\n\n    Arguments:\n        tensor (Tensor): tensor to split.\n        split_size (int): size of a single chunk.\n        dim (int): dimension along which to split the tensor.\n    \"\"\"\n    if dim < 0:\n        dim += tensor.dim()\n    dim_size = tensor.size(dim)\n    num_splits = (dim_size + split_size - 1) // split_size\n    last_split_size = split_size - (split_size * num_splits - dim_size)\n\n    def get_split_size(i):\n        return split_size if i < num_splits - 1 else last_split_size\n    return tuple(tensor.narrow(int(dim), int(i * split_size), int(get_split_size(i))) for i\n                 in _range(0, num_splits))\n\n\ndef chunk(tensor, chunks, dim=0):\n    \"\"\"Splits a tensor into a number of chunks along a given dimension.\n\n    Arguments:\n        tensor (Tensor): tensor to split.\n        chunks (int): number of chunks to return.\n        dim (int): dimension along which to split the tensor.\n    \"\"\"\n    if dim < 0:\n        dim += tensor.dim()\n    split_size = (tensor.size(dim) + chunks - 1) // chunks\n    return split(tensor, split_size, dim)\n\n\ndef stack(sequence, dim=0, out=None):\n    \"\"\"Concatenates sequence of tensors along a new dimension.\n\n    All tensors need to be of the same size.\n\n    Arguments:\n        sequence (Sequence): sequence of tensors to concatenate.\n        dim (int): dimension to insert. Has to be between 0 and the number\n            of dimensions of concatenated tensors (inclusive).\n    \"\"\"\n    if len(sequence) == 0:\n        raise ValueError(\"stack expects a non-empty sequence of tensors\")\n    if dim < 0:\n        dim += sequence[0].dim()\n    inputs = [t.unsqueeze(dim) for t in sequence]\n    if out is None:\n        return torch.cat(inputs, dim)\n    else:\n        return torch.cat(inputs, dim, out=out)\n\n\ndef unbind(tensor, dim=0):\n    \"\"\"Removes a tensor dimension.\n\n    Returns a tuple of all slices along a given dimension, already without it.\n\n    Arguments:\n        tensor (Tensor): tensor to unbind.\n        dim (int): dimension to remove.\n    \"\"\"\n    return tuple(tensor.select(dim, i) for i in _range(tensor.size(dim)))\n\n\ndef btriunpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True):\n    \"\"\"Unpacks the data and pivots from a batched LU factorization (btrifact) of a tensor.\n\n    Returns a tuple indexed by:\n      0: The pivots.\n      1: The L tensor.\n      2: The U tensor.\n\n    Arguments:\n        LU_data (Tensor): The packed LU factorization data.\n        LU_pivots (Tensor): The packed LU factorization pivots.\n        unpack_data (bool): Flag indicating if the data should be unpacked.\n        unpack_pivots (bool): Flag indicating if the pivots should be unpacked.\n    \"\"\"\n\n    nBatch, sz, _ = LU_data.size()\n\n    if unpack_data:\n        I_U = torch.triu(torch.ones(sz, sz)).type_as(LU_data).byte().unsqueeze(0).expand(nBatch, sz, sz)\n        I_L = 1 - I_U\n        L = LU_data.new(LU_data.size()).zero_()\n        U = LU_data.new(LU_data.size()).zero_()\n        I_diag = torch.eye(sz).type_as(LU_data).byte().unsqueeze(0).expand(nBatch, sz, sz)\n        L[I_diag] = 1.0\n        L[I_L] = LU_data[I_L]\n        U[I_U] = LU_data[I_U]\n    else:\n        L = U = None\n\n    if unpack_pivots:\n        P = torch.eye(sz).type_as(LU_data).unsqueeze(0).repeat(nBatch, 1, 1)\n        for i in range(nBatch):\n            for j in range(sz):\n                k = LU_pivots[i, j] - 1\n                t = P[i, :, j].clone()\n                P[i, :, j] = P[i, :, k]\n                P[i, :, k] = t\n    else:\n        P = None\n\n    return P, L, U\n\n\ndef matmul(tensor1, tensor2, out=None):\n    \"\"\"Matrix product of two tensors.\n\n    The behavior depends on the dimensionality of the tensors as follows:\n\n    - If both tensors are 1-dimensional, the dot product (scalar) is returned.\n    - If both arguments are 2-dimensional, the matrix-matrix product is returned.\n    - If the first argument is 1-dimensional and the second argument is 2-dimensional,\n      a 1 is prepended to its dimension for the purpose of the matrix multiply.\n      After the matrix multiply, the prepended dimension is removed.\n    - If the first argument is 2-dimensional and the second argument is 1-dimensional,\n      the matrix-vector product is returned.\n    - If both arguments are at least 1-dimensional and at least one argument is\n      N-dimensional (where N > 2), then a batched matrix multiply is returned.  If the first\n      argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the\n      batched matrix multiply and removed after.  If the second argument is 1-dimensional, a\n      1 is appended to its dimension for the purpose of the batched matrix multiple and removed after.\n      The non-matrix (i.e. batch) dimensions are :ref:`broadcasted <broadcasting-semantics>` (and thus\n      must be broadcastable).  For example, if :attr:`tensor1` is a `j x 1 x n x m` Tensor\n      and :attr:`tensor2` is a `k x m x p` Tensor, :attr:`out` will be an `j x k x n x p` Tensor.\n\n    .. note::\n\n        The 1-dimensional dot product version of this function does not support an :attr:`out` parameter.\n\n    Arguments:\n        tensor1 (Tensor): First tensor to be multiplied\n        tensor2 (Tensor): Second tensor to be multiplied\n        out (Tensor, optional): Output tensor\n    \"\"\"\n    dim_tensor1 = tensor1.dim()\n    dim_tensor2 = tensor2.dim()\n    if dim_tensor1 == 1 and dim_tensor2 == 1:\n        if out is None:\n            return torch.dot(tensor1, tensor2)\n        else:\n            raise ValueError(\"out must be None for 1-d tensor matmul, returns a scalar\")\n    if dim_tensor1 == 2 and dim_tensor2 == 1:\n        if out is None:\n            return torch.mv(tensor1, tensor2)\n        else:\n            return torch.mv(tensor1, tensor2, out=out)\n    elif dim_tensor1 == 1 and dim_tensor2 == 2:\n        if out is None:\n            return torch.mm(tensor1.unsqueeze(0), tensor2).squeeze_(0)\n        else:\n            return torch.mm(tensor1.unsqueeze(0), tensor2, out=out).squeeze_(0)\n    elif dim_tensor1 == 2 and dim_tensor2 == 2:\n        if out is None:\n            return torch.mm(tensor1, tensor2)\n        else:\n            return torch.mm(tensor1, tensor2, out=out)\n    elif dim_tensor1 >= 3 and (dim_tensor2 == 1 or dim_tensor2 == 2):\n        # optimization: use mm instead of bmm by folding tensor1's batch into\n        # its leading matrix dimension.\n\n        if dim_tensor2 == 1:\n            tensor2 = tensor2.unsqueeze(-1)\n\n        size1 = tensor1.size()\n        size2 = tensor2.size()\n        output_size = size1[:-1] + size2[-1:]\n\n        # fold the batch into the first dimension\n        tensor1 = tensor1.contiguous().view(-1, size1[-1])\n\n        if out is None or not out.is_contiguous():\n            output = torch.mm(tensor1, tensor2)\n        else:\n            output = torch.mm(tensor1, tensor2, out=out)\n\n        output = output.view(output_size)\n\n        if dim_tensor2 == 1:\n            output = output.squeeze(-1)\n\n        if out is not None:\n            out.set_(output)\n            return out\n\n        return output\n    elif (dim_tensor1 >= 1 and dim_tensor2 >= 1) and (dim_tensor1 >= 3 or dim_tensor2 >= 3):\n        # ensure each tensor size is at least 3-dimensional\n        tensor1_exp_size = torch.Size((1,) * max(3 - tensor1.dim(), 0) + tensor1.size())\n        # rhs needs to be a separate case since we can't freely expand 1s on the rhs, but can on lhs\n        if dim_tensor2 == 1:\n            tensor2 = tensor2.unsqueeze(1)\n        tensor2_exp_size = torch.Size((1,) * max(3 - tensor2.dim(), 0) + tensor2.size())\n\n        # expand the batch portion (i.e. cut off matrix dimensions and expand rest)\n        expand_batch_portion = torch._C._infer_size(tensor1_exp_size[:-2], tensor2_exp_size[:-2])\n\n        # flatten expanded batches\n        tensor1_expanded = tensor1.expand(*(expand_batch_portion + tensor1_exp_size[-2:])) \\\n            .contiguous().view(reduce(mul, expand_batch_portion), *tensor1_exp_size[-2:])\n        tensor2_expanded = tensor2.expand(*(expand_batch_portion + tensor2_exp_size[-2:])) \\\n            .contiguous().view(reduce(mul, expand_batch_portion), *tensor2_exp_size[-2:])\n\n        # reshape batches back into result\n        total_expansion = expand_batch_portion + (tensor1_exp_size[-2], tensor2_exp_size[-1])\n\n        def maybeSqueeze(tensor):\n            if dim_tensor1 == 1:\n                return tensor.squeeze(-2)\n            elif dim_tensor2 == 1:\n                return tensor.squeeze(-1)\n            else:\n                return tensor\n\n        if out is None or not out.is_contiguous():\n            output = torch.bmm(tensor1_expanded, tensor2_expanded)\n        else:\n            output = torch.bmm(tensor1_expanded, tensor2_expanded, out=out)\n\n        output = maybeSqueeze(output.view(total_expansion))\n\n        if out is not None:\n            out.set_(output)\n            return out\n\n        return output\n\n    raise ValueError(\"both arguments to __matmul__ need to be at least 1D, \"\n                     \"but they are {}D and {}D\".format(dim_tensor1, dim_tensor2))\n", "path": "torch/functional.py"}], "after_files": [{"content": "import torch\nfrom ._utils import _range\nfrom operator import mul\nfrom functools import reduce\n\n__all__ = [\n    'split', 'chunk', 'stack', 'unbind', 'btriunpack', 'matmul',\n]\n\n\ndef split(tensor, split_size, dim=0):\n    \"\"\"Splits the tensor into equally sized chunks (if possible).\n\n    Last chunk will be smaller if the tensor size along a given dimension\n    is not divisible by ``split_size``.\n\n    Arguments:\n        tensor (Tensor): tensor to split.\n        split_size (int): size of a single chunk.\n        dim (int): dimension along which to split the tensor.\n    \"\"\"\n    if dim < 0:\n        dim += tensor.dim()\n    dim_size = tensor.size(dim)\n    num_splits = (dim_size + split_size - 1) // split_size\n    last_split_size = split_size - (split_size * num_splits - dim_size)\n\n    def get_split_size(i):\n        return split_size if i < num_splits - 1 else last_split_size\n    return tuple(tensor.narrow(int(dim), int(i * split_size), int(get_split_size(i))) for i\n                 in _range(0, num_splits))\n\n\ndef chunk(tensor, chunks, dim=0):\n    \"\"\"Splits a tensor into a number of chunks along a given dimension.\n\n    Arguments:\n        tensor (Tensor): tensor to split.\n        chunks (int): number of chunks to return.\n        dim (int): dimension along which to split the tensor.\n    \"\"\"\n    if dim < 0:\n        dim += tensor.dim()\n    split_size = (tensor.size(dim) + chunks - 1) // chunks\n    return split(tensor, split_size, dim)\n\n\ndef stack(sequence, dim=0, out=None):\n    \"\"\"Concatenates sequence of tensors along a new dimension.\n\n    All tensors need to be of the same size.\n\n    Arguments:\n        sequence (Sequence): sequence of tensors to concatenate.\n        dim (int): dimension to insert. Has to be between 0 and the number\n            of dimensions of concatenated tensors (inclusive).\n    \"\"\"\n    if len(sequence) == 0:\n        raise ValueError(\"stack expects a non-empty sequence of tensors\")\n    if dim < 0:\n        dim += sequence[0].dim() + 1\n    inputs = [t.unsqueeze(dim) for t in sequence]\n    if out is None:\n        return torch.cat(inputs, dim)\n    else:\n        return torch.cat(inputs, dim, out=out)\n\n\ndef unbind(tensor, dim=0):\n    \"\"\"Removes a tensor dimension.\n\n    Returns a tuple of all slices along a given dimension, already without it.\n\n    Arguments:\n        tensor (Tensor): tensor to unbind.\n        dim (int): dimension to remove.\n    \"\"\"\n    return tuple(tensor.select(dim, i) for i in _range(tensor.size(dim)))\n\n\ndef btriunpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True):\n    \"\"\"Unpacks the data and pivots from a batched LU factorization (btrifact) of a tensor.\n\n    Returns a tuple indexed by:\n      0: The pivots.\n      1: The L tensor.\n      2: The U tensor.\n\n    Arguments:\n        LU_data (Tensor): The packed LU factorization data.\n        LU_pivots (Tensor): The packed LU factorization pivots.\n        unpack_data (bool): Flag indicating if the data should be unpacked.\n        unpack_pivots (bool): Flag indicating if the pivots should be unpacked.\n    \"\"\"\n\n    nBatch, sz, _ = LU_data.size()\n\n    if unpack_data:\n        I_U = torch.triu(torch.ones(sz, sz)).type_as(LU_data).byte().unsqueeze(0).expand(nBatch, sz, sz)\n        I_L = 1 - I_U\n        L = LU_data.new(LU_data.size()).zero_()\n        U = LU_data.new(LU_data.size()).zero_()\n        I_diag = torch.eye(sz).type_as(LU_data).byte().unsqueeze(0).expand(nBatch, sz, sz)\n        L[I_diag] = 1.0\n        L[I_L] = LU_data[I_L]\n        U[I_U] = LU_data[I_U]\n    else:\n        L = U = None\n\n    if unpack_pivots:\n        P = torch.eye(sz).type_as(LU_data).unsqueeze(0).repeat(nBatch, 1, 1)\n        for i in range(nBatch):\n            for j in range(sz):\n                k = LU_pivots[i, j] - 1\n                t = P[i, :, j].clone()\n                P[i, :, j] = P[i, :, k]\n                P[i, :, k] = t\n    else:\n        P = None\n\n    return P, L, U\n\n\ndef matmul(tensor1, tensor2, out=None):\n    \"\"\"Matrix product of two tensors.\n\n    The behavior depends on the dimensionality of the tensors as follows:\n\n    - If both tensors are 1-dimensional, the dot product (scalar) is returned.\n    - If both arguments are 2-dimensional, the matrix-matrix product is returned.\n    - If the first argument is 1-dimensional and the second argument is 2-dimensional,\n      a 1 is prepended to its dimension for the purpose of the matrix multiply.\n      After the matrix multiply, the prepended dimension is removed.\n    - If the first argument is 2-dimensional and the second argument is 1-dimensional,\n      the matrix-vector product is returned.\n    - If both arguments are at least 1-dimensional and at least one argument is\n      N-dimensional (where N > 2), then a batched matrix multiply is returned.  If the first\n      argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the\n      batched matrix multiply and removed after.  If the second argument is 1-dimensional, a\n      1 is appended to its dimension for the purpose of the batched matrix multiple and removed after.\n      The non-matrix (i.e. batch) dimensions are :ref:`broadcasted <broadcasting-semantics>` (and thus\n      must be broadcastable).  For example, if :attr:`tensor1` is a `j x 1 x n x m` Tensor\n      and :attr:`tensor2` is a `k x m x p` Tensor, :attr:`out` will be an `j x k x n x p` Tensor.\n\n    .. note::\n\n        The 1-dimensional dot product version of this function does not support an :attr:`out` parameter.\n\n    Arguments:\n        tensor1 (Tensor): First tensor to be multiplied\n        tensor2 (Tensor): Second tensor to be multiplied\n        out (Tensor, optional): Output tensor\n    \"\"\"\n    dim_tensor1 = tensor1.dim()\n    dim_tensor2 = tensor2.dim()\n    if dim_tensor1 == 1 and dim_tensor2 == 1:\n        if out is None:\n            return torch.dot(tensor1, tensor2)\n        else:\n            raise ValueError(\"out must be None for 1-d tensor matmul, returns a scalar\")\n    if dim_tensor1 == 2 and dim_tensor2 == 1:\n        if out is None:\n            return torch.mv(tensor1, tensor2)\n        else:\n            return torch.mv(tensor1, tensor2, out=out)\n    elif dim_tensor1 == 1 and dim_tensor2 == 2:\n        if out is None:\n            return torch.mm(tensor1.unsqueeze(0), tensor2).squeeze_(0)\n        else:\n            return torch.mm(tensor1.unsqueeze(0), tensor2, out=out).squeeze_(0)\n    elif dim_tensor1 == 2 and dim_tensor2 == 2:\n        if out is None:\n            return torch.mm(tensor1, tensor2)\n        else:\n            return torch.mm(tensor1, tensor2, out=out)\n    elif dim_tensor1 >= 3 and (dim_tensor2 == 1 or dim_tensor2 == 2):\n        # optimization: use mm instead of bmm by folding tensor1's batch into\n        # its leading matrix dimension.\n\n        if dim_tensor2 == 1:\n            tensor2 = tensor2.unsqueeze(-1)\n\n        size1 = tensor1.size()\n        size2 = tensor2.size()\n        output_size = size1[:-1] + size2[-1:]\n\n        # fold the batch into the first dimension\n        tensor1 = tensor1.contiguous().view(-1, size1[-1])\n\n        if out is None or not out.is_contiguous():\n            output = torch.mm(tensor1, tensor2)\n        else:\n            output = torch.mm(tensor1, tensor2, out=out)\n\n        output = output.view(output_size)\n\n        if dim_tensor2 == 1:\n            output = output.squeeze(-1)\n\n        if out is not None:\n            out.set_(output)\n            return out\n\n        return output\n    elif (dim_tensor1 >= 1 and dim_tensor2 >= 1) and (dim_tensor1 >= 3 or dim_tensor2 >= 3):\n        # ensure each tensor size is at least 3-dimensional\n        tensor1_exp_size = torch.Size((1,) * max(3 - tensor1.dim(), 0) + tensor1.size())\n        # rhs needs to be a separate case since we can't freely expand 1s on the rhs, but can on lhs\n        if dim_tensor2 == 1:\n            tensor2 = tensor2.unsqueeze(1)\n        tensor2_exp_size = torch.Size((1,) * max(3 - tensor2.dim(), 0) + tensor2.size())\n\n        # expand the batch portion (i.e. cut off matrix dimensions and expand rest)\n        expand_batch_portion = torch._C._infer_size(tensor1_exp_size[:-2], tensor2_exp_size[:-2])\n\n        # flatten expanded batches\n        tensor1_expanded = tensor1.expand(*(expand_batch_portion + tensor1_exp_size[-2:])) \\\n            .contiguous().view(reduce(mul, expand_batch_portion), *tensor1_exp_size[-2:])\n        tensor2_expanded = tensor2.expand(*(expand_batch_portion + tensor2_exp_size[-2:])) \\\n            .contiguous().view(reduce(mul, expand_batch_portion), *tensor2_exp_size[-2:])\n\n        # reshape batches back into result\n        total_expansion = expand_batch_portion + (tensor1_exp_size[-2], tensor2_exp_size[-1])\n\n        def maybeSqueeze(tensor):\n            if dim_tensor1 == 1:\n                return tensor.squeeze(-2)\n            elif dim_tensor2 == 1:\n                return tensor.squeeze(-1)\n            else:\n                return tensor\n\n        if out is None or not out.is_contiguous():\n            output = torch.bmm(tensor1_expanded, tensor2_expanded)\n        else:\n            output = torch.bmm(tensor1_expanded, tensor2_expanded, out=out)\n\n        output = maybeSqueeze(output.view(total_expansion))\n\n        if out is not None:\n            out.set_(output)\n            return out\n\n        return output\n\n    raise ValueError(\"both arguments to __matmul__ need to be at least 1D, \"\n                     \"but they are {}D and {}D\".format(dim_tensor1, dim_tensor2))\n", "path": "torch/functional.py"}]}