Add type annotation to PyTorch/XLA code and tests & expand tests for various types as needed

[miladm]

🐛 Bug

Most of the code in PyTorch/XLA does not have type annotations, and almost everything shows up as Any if you use a type checker. Let's fix that! Once we have good coverage, we can add a type checker to the CI to help us catch bugs.

Please file child bugs for specific files/modules as you go.

typing module reference: https://docs.python.org/3/library/typing.html

Please introduce a plan for the target code to be addressed by this bug and a corresponding set of tests to expand on.

[miladm]

@ysiraichi can you please take over this issue. @zpcore plz assist.