SpecDB: Add casting utilities

Summary:
Utility functions to return:
- all dtypes that a given dtype can be casted to
- all dtypes that can be casted to a given dtype

Reviewed By: SS-JIA

Differential Revision: D52919718

fbshipit-source-id: bb93b035a32df755f2f4879f99a411c95e04c56f

Author: manuelcandales

specdb/dtypes.py

@@ -14,4 +14,12 @@
 _floating_and_half = [torch.half] + _floating
 _complex = [torch.chalf, torch.cfloat, torch.cdouble]
 _quant = [torch.qint8, torch.quint8, torch.qint32, torch.quint4x2, torch.quint2x4]
-_all = [torch.bool] + _int + _floating + _complex + _quant
+_all = [torch.bool] + _int + _floating_and_half + _complex + _quant
+
+
+def can_cast_from(t):
+    return [x for x in _all if torch.can_cast(t, x)]
+
+
+def can_cast_to(t):
+    return [x for x in _all if torch.can_cast(x, t)]