Skip to content

feat: refactor dtypes as objects, add constants#11

Open
nicola-bastianello wants to merge 19 commits into
team-decent:mainfrom
nicola-bastianello:dtypes
Open

feat: refactor dtypes as objects, add constants#11
nicola-bastianello wants to merge 19 commits into
team-decent:mainfrom
nicola-bastianello:dtypes

Conversation

@nicola-bastianello

Copy link
Copy Markdown
Member

this PR refactors dtypes as objects; in particular:

  • it defines a new custom object dtype instantiated as dtype("float32"), with properties name, available and backend_dtype. backend_dtype is bound to the corresponding backend dtype object at set_backend, and any dtype created afterwards is also bound
  • two dtypes are the same if they are both available and have the same name
  • Backend now 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 of Backend return a native dtype object if available, None otherwise
  • functions involving dtypes were refactored to use dtype objects
  • to access the dtypes: from decent_array import float32
  • the user guide was updated to discuss dtypes; dtype object is documented but the dtype instances are not in the API reference

other changes:

  • added constants e, pi, inf, nan as aliases of the framework-native constants (after set_backend is called, before they alias the corresponding math constants). to get the constants: from decent_array import e
  • constants are also abstract properties of Backend that need to return the framework-native constants
  • renamed SupportedFrameworks, SupportedDevices, SupportedArrayTypes by removing "supported". I personally find them easier to read this way since the names start differently, but I'm open to discuss
  • added a private module _utils.py to store unwrap, which was redefined in several places

the dtype refactor has the following benefits:

  • it aligns with all frameworks, which also define custom dtype objects (and with the array api as well)
  • it allows to define dtypes/constants as properties of Backend, ensuring that we handle them in the same way as all other framework-native items
  • it allows to mark dtypes as available or not, since not all dtypes are available in all settings (factors that impact availability are framework, framework settings, OS, device). this can be expanded in the future to do programmatic checks of what is available in each framework (e.g. by running a test operation and catching errors; in most frameworks there is no other way to check availability)
  • it is more user friendly, as one just does dtype("float32")

closes #8

@Simpag Simpag left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

.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)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good idea!


@property
def bfloat16(self) -> torch.dtype:
return torch.bfloat16

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

when running on gpu, bfloat16 is only available on certain cards. There's a bfloat16 utility in torch

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll change to check explicitly


@property
def float64(self) -> torch.dtype:
return torch.float64

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

float64 is not supported on mps iirc, should be verified

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

indeed it is not

@nicola-bastianello

Copy link
Copy Markdown
Member Author

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

not sure I get what you mean, can you please clarify?

@nicola-bastianello

Copy link
Copy Markdown
Member Author

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

@nicola-bastianello

Copy link
Copy Markdown
Member Author

all fixed!

@Simpag if you clarify your comment below I'll look into that also

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Refactoring dtypes (and constants)

2 participants