I think this belongs to the IMPL function (it's neither a check nor output metadata setting).

Then, you would have to replicate it to the CUDA IMPL function, too.

The reason it's here is that the size of self is modified using expand_size (not in-place though, we create another tensor and then use it from there for the checks).

Nice job! It got brief and straight-forward.

I feel like we are losing this copy. But, somehow all the tests are passing.

I thought about it before, and I gave it a thought again.

My observation is, expand_size(self_, {batch1.size(0), batch1.size(1), batch2.size(2)}) calls self_.expand method (if required).

Now the thing is, if result tensor's size is already (or if we set explicitly using set_output): {batch1.size(0), batch1.size(1), batch2.size(2)} - and self is not of the required size, then copy_ will resize the self tensor to result tensor's size.

In [17]: M = torch.randn(10, 1, 5)
In [18]: _out = torch.empty((10, 3, 5))

In [20]: _out.copy_(M).size()
Out[20]: torch.Size([10, 3, 5])

And we do this copy_ in the impl function now (after your suggestion in the review comment above, earlier it was in the meta function) if result is not same as self.

Now, result will be same as self if:

1. It's an in-place call (auto-generated code sets result as self in the meta function itself).
2. out is self (like: torch.baddbmm(M, batch1, batch2, out=M)). -- I'm not resizing for this case, and that's where it's wrong, please see below

On upstream:

In [43]: M = torch.randn(10, 1, 5)

In [44]: torch.baddbmm(M, batch1, batch2, out=M).size()
Out[44]: torch.Size([10, 3, 5])

On this branch:

In [8]: M = torch.randn(10, 1, 5)

In [9]: torch.baddbmm(M, batch1, batch2, out=M).size()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-9-4d21e89d2e08> in <module>
----> 1 torch.baddbmm(M, batch1, batch2, out=M).size()

RuntimeError: Expected self_sizes[0] == bs && self_sizes[1] == res_rows && self_sizes[2] == res_cols to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)

Now, interesting case is for in-place, this copy and expansion of size doesn't happen in upstream (rightly so). The problem is to find a way if the meta function is called from in-place dispatch, or method/function dispatch. Not sure how to do this :/ Will have to ponder on this one.

In [42]: M.baddbmm_(batch1, batch2).size()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-42-ffc510337fae> in <module>
----> 1 M.baddbmm_(batch1, batch2).size()

RuntimeError: Expected self_sizes[0] == bs && self_sizes[1] == res_rows && self_sizes[2] == res_cols to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)

We need to test this behavior as well in test_torch.py. I'll try to send a fix for this one soon.

Thanks for asking this question, @ysiraichi - sorry for the long message. ;)

I'm stuck here:

How would you differentiate these 2 calls (see below) from meta function? (Note: please see the comment above on why we need to differentiate b/w these calls)

>>> M.baddbmm_(batch1, batch2)  # [1] In-Place
>>> torch.baddbmm(M, batch1, batch2, out=M)  # [2] out=self tensor

From the auto-generated code for in-place variant in build/aten/src/ATen/RegisterCPU.cpp:

struct structured_baddbmm_out_cpu_inplace final : public at::native::structured_baddbmm_out_cpu {
     structured_baddbmm_out_cpu_inplace(Tensor& self) : outputs_{std::ref(self)} {}
     // ...
 };

As outputs_ is initialized with std::ref(self) so result.is_same(self) is always gonna be true.

cc: @ezyang @bdhirsh @ysiraichi

Similarly for [2] out=M call, result.is_same(self) will be true. How should we differentiate b/w these calls?

Hmm, this actually seems more in-line with our out= behavior: for most other out= ops, if you pass in an output tensor with an incorrect size we raise an error instead of transparently resizing the output for you (unless the tensor has zero size).

I'm just not sure if we want a deprecation cycle, since it's not clear to me how commonly this behavior is relied on (if I had to guess though, it seems like a pretty uncommon use case?). I also don't think you have an easy way with structured kernels to differentiate those two cases.

Now, I encountered this behaviour in a test in test/test_linalg.py the other day:

This will be fixed in #64387 , but it signals that there may be other places where this behaviour is used.

Note that, at the moment, if an incorrectly shaped out is used, a warning is issued. As such, I think that just making that warning into an error may be a reasonable behaviour here.

Thanks for the valuable inputs, @lezcano, and @bdhirsh. Following the suggestions, we now raise an error for incorrectly shaped out (earlier we allowed resizing). However, this behavior is only for baddbmm. We resize the output for bmm, there is a test that uses incorrect shapes for out here: https://github.com/pytorch/pytorch/blob/master/test/test_torch.py#L8394.

The only thing that concerns me is the error messages for inplace call and out=M call. The error with in-place call should have self_sizes mentioned instead of result_sizes, but that will need distinction b/w the 2 calls.

>>> M = torch.randn(10, 1, 5)  # Incorrect shape
>>> batch1 = torch.randn(10, 3, 4)
>>> batch2 = torch.randn(10, 4, 5)

>>> M.baddbmm_(batch1, batch2)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Expected result_sizes[0] == bs && result_sizes[1] == res_rows && result_sizes[2] == res_cols to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)

>>> torch.baddbmm(M, batch1, batch2, out=M)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Expected result_sizes[0] == bs && result_sizes[1] == res_rows && result_sizes[2] == res_cols to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)

The tests might need another look, which I personally feel can be done in a follow-up PR.

How would you differentiate these 2 calls (see below) from meta function? (Note: please see the comment above on why we need to differentiate b/w these calls)

While I agree with @bdhirsh (i.e. you can't), you can resize the output tensors conditionally by using set_output. So, instead of writing:

auto& result = meta.maybe_get_output(0);
if (is_bmm || !result.defined()) {
  meta.set_output({bs, res_rows, res_cols}, batch1.options());
} else {
  const auto result_sizes = result.sizes();
  TORCH_CHECK(result_sizes[0] == bs && result_sizes[1] == res_rows && result_sizes[2] == res_cols);
}

You could write:

// 'set_output' does not resize for in-place calls
meta.set_output({bs, res_rows, res_cols}, batch1.options());
// Error is raised if called from the in-place overload with incorrect self shape
TORCH_CHECK(result_sizes[0] == bs && result_sizes[1] == res_rows && result_sizes[2] == res_cols);

@ezyang @bdhirsh @lezcano @krshrimali
As far as I remember, the idea was to turn the warning into an error (as @lezcano pointed out) in the future, by using resize_output function (for consistency). Since that function is called by set_output, maybe we could rely on that for now (instead of directly raising an error for the out-of-place). Not sure, though. What do you think?

@ysiraichi good catch, that seems like it should work (and also raise the warning). It's probably worth quickly confirming that with that code change, we now get the deprecation warning:

>>> M = torch.randn(10, 1, 5)  # Incorrect shape
>>> batch1 = torch.randn(10, 3, 4)
>>> batch2 = torch.randn(10, 4, 5)

>>> torch.baddbmm(M, batch1, batch2, out=M)
# (we should see a deprecation warning, because meta.set_output calls at::native::set_output)

This is really cool, @ysiraichi! I wasn't aware that meta.set_output() doesn't resize for in-place calls. This helps!

Here is how the use-cases look like with the revised changes: (cc: @bdhirsh)

In [11]: M = torch.randn(10, 1, 5)
In [12]: batch1 = torch.randn(10, 3, 4)
In [13]: batch2 = torch.randn(10, 4, 5)

In [14]: torch.baddbmm(M, batch1, batch2).size()
Out[14]: torch.Size([10, 3, 5])

In [15]: M.baddbmm_(batch1, batch2).size()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-15-ffc510337fae> in <module>
----> 1 M.baddbmm_(batch1, batch2).size()

RuntimeError: Expected result_sizes[0] == bs && result_sizes[1] == res_rows && result_sizes[2] == res_cols to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)

In [16]: torch.baddbmm(M, batch1, batch2, out=M).size()
<ipython-input-16-4d21e89d2e08>:1: UserWarning: An output with one or more elements was resized since it had shape [10, 1, 5], which does not match the required output shape [10, 3, 5].This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at  ../aten/src/ATen/native/Resize.cpp:16.)
  torch.baddbmm(M, batch1, batch2, out=M).size()
Out[16]: torch.Size([10, 3, 5])

We do see a deprecation warning with out=M, and we also get an error for the in-place call. I've made the required commits, the tests should pass and this PR should hopefully look good now. Welcoming reviews now. :)

Nit: an error message here would be nice!

You don't have to spell it out individually, ArrayRef has equality defined. This also solves a bug where you can OOB access if the dimensionality differs. Also another one for #57827

@ysiraichi maybe you can fix the underlying issue, so we don't have to keep pasting this in? :)

Right, to me figuring out the error message gets tricky here (but I've still made an effort to be as correct as possible, in the recent commit) because an error is raised if:

1. out tensor is passed with an incorrect shape.
2. (in-place call), input tensor is passed with an incorrect shape.

I've added an error message in the recent commit, please let me know if it sounds good or if you have any suggestions. However, we'll not reach this error frequently, a lot of times it will fail at expand_size:

In [8]: M = torch.randn(1, 3, 2)

In [9]: M.baddbmm_(batch1, batch2)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-9-a9e665d29b16> in <module>
----> 1 M.baddbmm_(batch1, batch2)

RuntimeError: The expanded size of the tensor (5) must match the existing size (2) at non-singleton dimension 2.  Target sizes: [10, 3, 5].  Tensor sizes: [1, 3, 2]

The only time when we'll get this error is for in-place calls, as expected:

In [12]: M = torch.randn(10, 1, 5)

# batch1 size: (10, 3, 4) and batch2 size: (10, 4, 5)
In [13]: M.baddbmm_(batch1, batch2).size()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-13-ffc510337fae> in <module>
----> 1 M.baddbmm_(batch1, batch2).size()

RuntimeError: Expected an output tensor with shape [10, 3, 5] but got shape [10, 1, 5]

Hope this error message looks good now. For completeness, I've also added a similar error message for self_sizes check in the same function. :)

You don't have to spell it out individually, ArrayRef has equality defined. This also solves a bug where you can OOB access if the dimensionality differs. Also another one for #57827

Sorry, I just noticed this. This should be fixed now, defined output_size (IntArrayRef) since it's used twice.

Also added error messages here:

# batch2 is of incorrect size
In [7]: M.baddbmm_(batch1, batch2)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-7-a9e665d29b16> in <module>
----> 1 M.baddbmm_(batch1, batch2)

RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [10, 4] but got: [10, 5].

Relevant commit

cc: @ezyang @ysiraichi