Upstream `argmax` shape function.
Keeping this first commit simple to test out the flow. Will bulk-add the
rest once this one goes through.
Shape function taken from:
https://github.com/llvm/torch-mlir/blob/5192a4e9f3e4cbc77838b70153cd4632aa43dd7f/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py#L488
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76592
Approved by: https://github.com/eellison
