diff --git a/tensor-shapes/pyrefly-torch-stubs/test/test_dtype_aliases.py b/tensor-shapes/pyrefly-torch-stubs/test/test_dtype_aliases.py index 7ebf9f256d..9e3a390d1c 100644 --- a/tensor-shapes/pyrefly-torch-stubs/test/test_dtype_aliases.py +++ b/tensor-shapes/pyrefly-torch-stubs/test/test_dtype_aliases.py @@ -15,3 +15,8 @@ def test_float_dtype_alias() -> None: assert torch.float is torch.float32 assert isinstance(torch.float, torch.dtype) assert_shape(torch.ones((2, 3), dtype=torch.float).shape, (2, 3)) + + +def test_half_dtype_alias() -> None: + assert torch.half is torch.float16 + assert_shape(torch.ones((2, 3), dtype=torch.half).shape, (2, 3)) diff --git a/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi b/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi index 32a3485e8f..7a4d55c6fd 100644 --- a/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi +++ b/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi @@ -695,6 +695,7 @@ type _LegacyTensorScalar = builtins.bool | builtins.int | builtins.float qint8: Any quint8: Any float16: Any +half: Any float32: Any float64: Any int8: Any