[WIP] Add autocast in torchax

[qihqi]

How torch.autocast works:

When you do

with torch.autocast(device):
  some math

torch will do the following:

1. Check the device module (in our case it's in the device_module.py file) which is registered here: https://github.com/pytorch/xla/blob/master/torchax/torchax/__init__.py#L92; on what dtypes are supported for autocast.
2. switch dispatch key for the underlying math. For example, CPU device -> AutocastCPU; CUDA -> AutocastCUDA. We use the PrivateUse1 dispatch key so this will dispatch with AutocastPrivateUse1. Usual case of CUDA, it will call an op registered to it: for example, for AutocastCPU, it is registered here: https://github.com/pytorch/pytorch/blob/05faba40287cf7d8734da96cb2e904f39710bf29/aten/src/ATen/autocast_mode.cpp#L322
   a. Above, L332 registers fallback meaning every op to AutocastCPU will automatically fallback to next (presumably to CPU), i.e. unless otherwise specified, it will be no-op. This line prevents errors. and then lines 329 and beyond defines those op that do have additional behavior; these are generated by a cpp macro + templates.

What we need to do to support autocast:

1. add this function

def get_amp_supported_dtype():
  return [torch.float16, torch.bfloat16]

to tell torch that we do support autocast. At this point it will run but will not actually do autocast yet

2. Change tensor base device to 'privateuseone' instead of meta in here: https://github.com/pytorch/xla/blob/master/torchax/torchax/tensor.py#L62 because if we use meta device we won't get dispatch from AutocastPrivateUse1. At this point we try running autocast we would see errors on op not registered for AutocastPrivateUse1.
3. Now we want to register the ops, and we can do that in Python using the torch.library API. The exact incantation is in the autocat_policies.py file. In that file, we still need to reimplement the logic of autocast, mainly downcast the input before calling certain ops.
4. Fix up errors from changing device from meta to privateuseone introduced in 2. There are many decompositions that calls into CPP and it will check Cpp registrations. Device Guard is one of those: (note: filed Ability to set device guard in Python pytorch#156052 to do that in python). This can be fixed with adding a Cpp file that calls the Guard registration (we don't really need to use the features provided by the device guard, just need to register it so Pytorch doesn't complain).
5. There is one more tests fail (on autograd from j2t_autograd) likely more cpp registration is needed.
6. Move the building of cpp file from runtime to compile time by fiddling with pyproject.toml / setup.py etc.