Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Draft
wants to merge 12 commits into
base: main
Choose a base branch
from

Conversation

havogt
Copy link
Contributor

@havogt havogt commented Dec 4, 2024

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 Outdated Show resolved Hide resolved
allocator: FieldBufferAllocationUtil = actual_allocator,
device: core_defs.Device = device,
) -> core_defs.NDArrayObject:
# TODO check how to get from FieldBufferAllocationUtil to FieldBufferAllocatorProtocol
Copy link
Contributor Author

Choose a reason for hiding this comment

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

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

@havogt
Copy link
Contributor Author

havogt commented Dec 6, 2024

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

@havogt havogt changed the title feat[next]: allocator in field feat[next]: Support for array api namespace as allocator Dec 20, 2024
@havogt 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
self.ndarray,
domain=self.domain,
dtype=self.dtype,
copy=True, # aligned_index???
Copy link
Contributor Author

Choose a reason for hiding this comment

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

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

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

Choose a reason for hiding this comment

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

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

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.

1 participant