[MLIR][TORCH] Add op verifier for aten.index_put op

[vivekkhandelwal1]

This commit adds an op verifier for the index_put op.

[sogartar]

I was expecting that the op would get a verifier that can be ran at any point. For example this can not be used to verify input IR that contains said op.
What is the reasoning to do the check during conversion and not have a verifier?

Also it is desirable to have a test to show that the check works as expected.

[sogartar]

You mentioned that for dynamic shapes the lowering would insert runtime asserts. How does this work?

[sogartar]

Shouldn't the verification fail if does not have sizes. I assume this is the shape. What does it mean even if there is no size?

[sogartar]

Shouldn't we fail here as the other case?

[sogartar]

Shouldn't we fail here as the other case?

[sogartar]

The indices may not be broadcast compatible on their own. E.g.

t = torch.empty([3, 4, 5], dtype=torch.int32)
indices = [
    torch.zeros(size=[5], dtype=torch.long),
    torch.zeros(size=[3], dtype=torch.long),
]
t[indices]

raises

    t[indices]
    ~^^^^^^^^^
IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [5], [3]

This function should be failable.

The resulting shape after indexing is more complicated. I think this verification would reject valid IR.
Example:

indices = [
    torch.zeros(size=[7], dtype=torch.long),
]
print(t[indices].shape)

torch.Size([7, 4, 5])

---

indices = [
    torch.zeros(size=[7], dtype=torch.long),
    torch.zeros(size=[7], dtype=torch.long),
]
print(t[indices].shape)

torch.Size([7, 5])

---

indices = [
    torch.zeros(size=[7], dtype=torch.long),
    torch.zeros(size=[7], dtype=torch.long),
    torch.zeros(size=[7], dtype=torch.long),
]
print(t[indices].shape)

torch.Size([7])

---

indices = [
    torch.zeros(size=[7], dtype=torch.long),
    torch.zeros(size=[7], dtype=torch.long),
    torch.zeros(size=[7], dtype=torch.long),
    torch.zeros(size=[7], dtype=torch.long),
]
print(t[indices].shape)

    print(t[indices].shape)
          ~^^^^^^^^^
IndexError: too many indices for tensor of dimension 3

Here I am using t[indices] indexing because according to the doc torch.Tensor.index_put_ is equivalent to tensor[indices] = values.

[sogartar]

I am actually surprised that there are no functions already available in Torch MLIR to get the shape after broadcasting or to check if broadcastable.