simplify boundingbox kernels with jit compatible ellipsis

[pmeier]

Addresses #5781 (comment). Cross posting my answer there:

Good catch. I was under the impression ... was not supported by torch.jit.script in general, but this seems fine:

@torch.jit.script
def foo(x: torch.Tensor) -> torch.Tensor:
    y = x.clone()
    y[..., 0::2] += 1
    return y

It seems, only the explicit indices are not supported:

@torch.jit.script
def bar(x: torch.Tensor) -> torch.Tensor:
    y = x.clone()
    y[..., [0, 2]] += 1
    return y

RuntimeError: 
Ellipses followed by tensor indexing is currently not supported:
[...]
def bar(x: torch.Tensor) -> torch.Tensor:
    y = x.clone()
    y[..., [0, 2]] += 1
    ~~~~~~~~~~~~~ <--- HERE
    return y

[datumbox]

LGTM.

If you look into the codebase for .view(-1, 4) you will find multiple other occurrences (for example at horizontal_flip_bounding_box, vertical_flip_bounding_box, affine_bounding_box etc. Should we fix this as well?

@vfdev-5 thoughts?

[pmeier]

The examples you mentioned are all using explicit indices and thus cannot simplified as explained above:

vision/torchvision/prototype/transforms/functional/_geometry.py

Line 31 in a64c674

bounding_box[:, [0, 2]] = image_size[1] - bounding_box[:, [2, 0]]

vision/torchvision/prototype/transforms/functional/_geometry.py

Line 97 in a64c674

bounding_box[:, [1, 3]] = image_size[0] - bounding_box[:, [3, 1]]

vision/torchvision/prototype/transforms/functional/_geometry.py

Line 239 in a64c674

points = bounding_box[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].view(-1, 2)

[vfdev-5]

IMO, view should be very tiny op. Great that we could simplify the code for this case and I think it is not a big deal that we couldn't for other cases... Maybe later when JIT will support that.
I think we can merge this PR.