feat[next]: Support for Array Api namespace as allocator #1771

havogt · 2024-12-04T17:14:51Z

Allow Array API namespaces as allocators in gtx.constructors. This allows e.g. to construct jax fields in non-hacky way.

Additional:

enable JAX embedded execution
add __copy__, __deepcopy__ to NDArrayField (with same memory layout as source Field)

src/gt4py/next/constructors.py

havogt · 2024-12-05T12:06:04Z

src/gt4py/next/allocators.py

+        allocator: FieldBufferAllocationUtil = actual_allocator,
+        device: core_defs.Device = device,
+    ) -> core_defs.NDArrayObject:
+        # TODO check how to get from FieldBufferAllocationUtil to FieldBufferAllocatorProtocol


think about long names -> maybe rename to NDArray...

__copy__, __deepcopy__

as_ndarray(allocator: ConcreteAllocator, copy: Optional[bool])

can TensorBuffer be removed?

havogt · 2024-12-06T15:37:11Z

think about long names -> maybe rename to NDArray...
__copy__, __deepcopy__
as_ndarray(allocator: ConcreteAllocator, copy: Optional[bool])
can TensorBuffer be removed?

havogt · 2024-12-20T09:28:36Z

src/gt4py/next/embedded/nd_array_field.py

+            self.ndarray,
+            domain=self.domain,
+            dtype=self.dtype,
+            copy=True,  # aligned_index???


Actually this is a missing piece in the current implementation. We should probably provide this as a default in all functions.

havogt · 2024-12-20T09:30:48Z

src/gt4py/next/_allocators.py

+def get_array_allocation_namespace(
+    allocator: Optional[FieldBufferAllocationUtil | core_defs.ArrayApiNamespace],
+    device: Optional[core_defs.Device] = None,
+) -> GTArrayAllocationNamespace:


Maybe put aligned_index here and add it as a default to the construction functions.
Same for device.

egparedes

I added some basic comments, but it's not a detailed review. I agree with should use this strategy short-term to be able to use jax and similar immutable data array api libraries, but I'm not sure yet if the current implementation strategy is good enough or it should be refactored. I'd propose to discuss it further offline.

egparedes · 2025-01-08T17:12:09Z

src/gt4py/_core/definitions.py

+    def empty(self, shape: Sequence[int], *, dtype: Any = None, device: Any = None) -> Any: ...
+    def zeros(self, shape: Sequence[int], *, dtype: Any = None, device: Any = None) -> Any: ...
+    def ones(self, shape: Sequence[int], *, dtype: Any = None, device: Any = None) -> Any: ...
+    def full(
+        self, shape: Sequence[int], fill_value: Scalar, *, dtype: Any = None, device: Any = None
+    ) -> Any: ...
+    def asarray(self, obj: Any, *, dtype: Any = None, copy: Any = None) -> Any: ...


I guess using NDArrayObject as output type doesn't work with mypy, right?

egparedes · 2025-01-08T17:15:30Z

src/gt4py/_core/definitions.py

+    # TODO(havogt): add relevant methods and attributes or wait for the standard to provide it, see e.g. https://github.com/data-apis/array-api/issues/697
+
+
+def is_array_api_namespace(obj: Any) -> TypeGuard[ArrayApiNamespace]:


We should start using TypeIs (https://devdocs.io/python~3.13/library/typing#typing.TypeIs) instead of TypeGuard whenever possible.

Suggested change

def is_array_api_namespace(obj: Any) -> TypeGuard[ArrayApiNamespace]:

def is_array_api_namespace(obj: Any) -> TypeIs[ArrayApiNamespace]:

egparedes · 2025-01-09T15:39:18Z

src/gt4py/next/_allocators.py

-    if device is None and allocator is None:
-        raise ValueError("No 'device' or 'allocator' specified.")
-    actual_allocator = get_allocator(allocator)
-    if actual_allocator is None:
-        assert device is not None  # for mypy
-        actual_allocator = device_allocators[device.device_type]
-    elif device is None:
-        device = core_defs.Device(actual_allocator.__gt_device_type__, 0)
-    elif device.device_type != actual_allocator.__gt_device_type__:
-        raise ValueError(f"Device '{device}' and allocator '{actual_allocator}' are incompatible.")


If this part is gone, the docstring needs to be updated, since most of the parameters are now compulsory.

egparedes · 2025-01-09T15:40:12Z

src/gt4py/next/_allocators.py

@@ -41,6 +43,63 @@
 )


+class GTArrayAllocationNamespace(Protocol):


I don't really like this name but haven't found a better one yet

egparedes · 2025-01-09T15:41:13Z

src/gt4py/next/_allocators.py

+    device: Optional[core_defs.Device], aligned_index: Optional[Sequence[common.NamedIndex]]
+) -> None:
+    if aligned_index is not None:
+        raise NotImplementedError("Aligned index is not support for Array API namespaces.")


Suggested change

raise NotImplementedError("Aligned index is not support for Array API namespaces.")

raise NotImplementedError("Aligned index is not supported for Array API namespaces.")

egparedes · 2025-01-09T15:47:50Z

tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py

+def test_copy(copy, nd_array_implementation):
+    testee = _make_field_or_scalar([[0, 1], [2, 3]], nd_array_implementation)
+    result = copy(testee)
+    assert np.array_equal(testee.ndarray, result.ndarray)


I would also test that testee and result (and their ndarrays) are not the same object

egparedes · 2025-01-10T08:29:09Z

tests/next_tests/unit_tests/test_constructors.py

+    return request.param
+
+
+def test_empty(allocator_device_refnamespace):


Just a suggestion: we could avoid the intermediate fixture here and just parametrize the test explicitly:

Suggested change

def test_empty(allocator_device_refnamespace):

@pytest.mark.parametrize("allocator, device, xp", [*allocator_device_refnamespace_params()])

def test_empty(allocator, device, xp):

egparedes · 2025-01-10T08:32:27Z

tests/next_tests/unit_tests/test_constructors.py


-    assert np.array_equal(a.ndarray, ref)
+
+def test_copy():


Isn't this test the same than the one in the embedded test_nd_array_field tests?

feat[next]: allocator in field

81d5cda

havogt commented Dec 5, 2024

View reviewed changes

src/gt4py/next/constructors.py Outdated Show resolved Hide resolved

havogt commented Dec 5, 2024

View reviewed changes

havogt added 11 commits December 9, 2024 08:58

continue refactoring

fe09c5f

save state

e4dabc8

cleanup copying

1b3e0c5

fix storage inferface test

afbbc8e

move to_array_api_dtype

3bb206f

fix variable names

1c5acd6

refactor to an allocation namespace

92b8ea6

fix compat

7817ec3

fix tests

6709b04

move device to allocation namespace construction

3ee60b8

fix get_actual_allocator

4fb8bab

havogt changed the title ~~feat[next]: allocator in field~~ feat[next]: Support for array api namespace as allocator Dec 20, 2024

havogt changed the title ~~feat[next]: Support for array api namespace as allocator~~ feat[next]: Support for Array Api namespace as allocator Dec 20, 2024

havogt commented Dec 20, 2024

View reviewed changes

egparedes reviewed Jan 10, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat[next]: Support for Array Api namespace as allocator #1771

feat[next]: Support for Array Api namespace as allocator #1771

havogt commented Dec 4, 2024 •

edited

Loading

havogt Dec 5, 2024

havogt commented Dec 6, 2024 •

edited

Loading

havogt Dec 20, 2024

havogt Dec 20, 2024

egparedes left a comment

egparedes Jan 8, 2025

egparedes Jan 8, 2025

egparedes Jan 9, 2025

egparedes Jan 9, 2025

egparedes Jan 9, 2025

egparedes Jan 9, 2025

egparedes Jan 10, 2025

egparedes Jan 10, 2025

		# TODO(havogt): add relevant methods and attributes or wait for the standard to provide it, see e.g. https://github.com/data-apis/array-api/issues/697


		def is_array_api_namespace(obj: Any) -> TypeGuard[ArrayApiNamespace]:

		@@ -41,6 +43,63 @@
		)


		class GTArrayAllocationNamespace(Protocol):

	raise NotImplementedError("Aligned index is not support for Array API namespaces.")
	raise NotImplementedError("Aligned index is not supported for Array API namespaces.")

		return request.param


		def test_empty(allocator_device_refnamespace):

	def test_empty(allocator_device_refnamespace):
	@pytest.mark.parametrize("allocator, device, xp", [*allocator_device_refnamespace_params()])
	def test_empty(allocator, device, xp):

feat[next]: Support for Array Api namespace as allocator #1771

Are you sure you want to change the base?

feat[next]: Support for Array Api namespace as allocator #1771

Conversation

havogt commented Dec 4, 2024 • edited Loading

Choose a reason for hiding this comment

havogt commented Dec 6, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

egparedes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

havogt commented Dec 4, 2024 •

edited

Loading

havogt commented Dec 6, 2024 •

edited

Loading