feat: refactor dtypes as objects, add constants#11
Conversation
Simpag
left a comment
There was a problem hiding this comment.
Overall looks fine, just some minor things. Would like to see some tests that validates each dtype for common devices (cpu, gpu and mps on the mac test env). Make sure that the expected exceptions are raise so we can catch any future changes to support
| raise ValueError(f"Unsupported dtype '{dtype.value}' for NumPy backend.") | ||
| return Array(jnp.asarray(x.value, dtype=_DTYPE_MAP[dtype])) | ||
| def astype(self, x: Array, dtype: dtype) -> Array: | ||
| if dtype not in _ALL_DTYPES.values(): |
There was a problem hiding this comment.
.values() will be an O(N) operation. Might be good to cache the available dtypes something like self._dtypes = set(_ALL_DTYPES.values()) for O(1)
|
|
||
| @property | ||
| def bfloat16(self) -> torch.dtype: | ||
| return torch.bfloat16 |
There was a problem hiding this comment.
when running on gpu, bfloat16 is only available on certain cards. There's a bfloat16 utility in torch
There was a problem hiding this comment.
I'll change to check explicitly
|
|
||
| @property | ||
| def float64(self) -> torch.dtype: | ||
| return torch.float64 |
There was a problem hiding this comment.
float64 is not supported on mps iirc, should be verified
There was a problem hiding this comment.
indeed it is not
not sure I get what you mean, can you please clarify? |
|
I don't understand why mypy-related stuff fails, it's all green on my computer. fixing the version doesn't fix, so it's something else. I'll look into this when I have time |
|
all fixed! @Simpag if you clarify your comment below I'll look into that also
|
this PR refactors dtypes as objects; in particular:
dtypeinstantiated asdtype("float32"), with propertiesname,availableandbackend_dtype.backend_dtypeis bound to the corresponding backend dtype object atset_backend, and any dtype created afterwards is also boundBackendnow exposes one property for each of the supported dtypes; the property is abstract if the dtype is supported by all frameworks, non-abstract and returning None otherwise. subclasses ofBackendreturn a native dtype object if available, None otherwisedtypeobjectsfrom decent_array import float32dtypeobject is documented but the dtype instances are not in the API referenceother changes:
e,pi,inf,nanas aliases of the framework-native constants (afterset_backendis called, before they alias the correspondingmathconstants). to get the constants:from decent_array import eBackendthat need to return the framework-native constantsSupportedFrameworks,SupportedDevices,SupportedArrayTypesby removing "supported". I personally find them easier to read this way since the names start differently, but I'm open to discuss_utils.pyto storeunwrap, which was redefined in several placesthe dtype refactor has the following benefits:
Backend, ensuring that we handle them in the same way as all other framework-native itemsdtype("float32")closes #8