Yeounoh Chung e2e9d15726 Unblock float16 dtype for xla autocasting (#109554) ...

`torch.autocast` with `xla` backend has been restricted to `torch.bfloat16`. This shouldn't be the case anymore.

This works with `xla::cast( ..., type=f16)`
```
IR {
  %0 = f32[] prim::Constant(), xla_shape=f32[], value=1
  %1 = f32[3,2]{1,0} aten::expand(%0), xla_shape=f32[3,2]{1,0}, size=(3, 2), dynamic_dims=(0, 0)
  %2 = f16[3,2]{1,0} xla::cast(%1), xla_shape=f16[3,2]{1,0}, type=f16, dtype=Half, stype=Float
  %3 = f32[] prim::Constant(), xla_shape=f32[], value=1
  %4 = f32[2,3]{1,0} aten::expand(%3), xla_shape=f32[2,3]{1,0}, size=(2, 3), dynamic_dims=(0, 0)
  %5 = f16[2,3]{1,0} xla::cast(%4), xla_shape=f16[2,3]{1,0}, type=f16, dtype=Half, stype=Float
  %6 = f16[2,2]{1,0} aten::mm(%5, %2), xla_shape=f16[2,2]{1,0}, ROOT=0
}
```

This will allow PyTorch/XLA to extend its autocast implementation to use `xla` backend for `float16` type as well.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/109554
Approved by: https://github.com/JackCaoG, https://github.com/bdhirsh

2023-09-21 03:19:44 +00:00

..

__init__.py

Apply UFMT to low traffic torch modules (#106249)

2023-07-29 23:37:30 +00:00

autocast_mode.py

Unblock float16 dtype for xla autocasting (#109554)

2023-09-21 03:19:44 +00:00