Thanks! Can you add test for U16, F64 etc on TPU CI to make sure they are actually natively supported?

I tested this manually, but I'll make sure there's a unit test for each one.

Thanks! Can you add test for U16, F64 etc on TPU CI to make sure they are actually natively supported?

Added a test that iterates through the dtypes that both torch and XLA support.

I found the list of supported dtypes on TPU internally, which does not include complex128, so using that dtype on TPU will trigger an error. The others are all supported now. We have two options for complex128:

1. Change our behavior to throw an error for unsupported dtypes. IMO this is the safer behavior. Plus, SPMD obscures the actual device type, which makes it hard to determine what the "real" device is without accessing runtime state.
2. Add downcasting back for complex128 only.

yea.. I don't know how many people actually uses complex128, is it supported on GPU? As we are increasing the effort on GPU this year, I don't want to regress PyTorch/XLA:GPU

yea.. I don't know how many people actually uses complex128, is it supported on GPU? As we are increasing the effort on GPU this year, I don't want to regress PyTorch/XLA:GPU

This PR won't affect behavior on GPU. I'm basically just removing the downcast when creating tensors on TPUs, and we never changed the dtype on GPU. If the underlying XLA:GPU client doesn't support complex128 (it probably does), then it would just throw an error.

Hmm, both CPU and GPU fail the test I added with complex128:

Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/absl/testing/parameterized.py", line 323, in bound_param_test
    return test_method(self, testcase_params)
  File "/tmp/pytorch/xla/test/pjrt/test_dtypes.py", line 22, in test_float_round_trip
    torch.testing.assert_close(xt.cpu(), t)
  File "/opt/conda/lib/python3.8/site-packages/torch/testing/_comparison.py", line 1520, in assert_close
    raise error_metas[0].to_error(msg)
AssertionError: Tensor-likes are not close!

Mismatched elements: 9 / 9 (100.0%)
Greatest absolute difference: 1.3942428785896892 at index (2, 1) (up to 1e-07 allowed)
Greatest relative difference: 0.9994087028336766 at index (1, 0) (up to 1e-07 allowed)

I'll just remove this test case for now.