parkipy.utils.get_array_module

parkipy.utils.get_array_module(

execution_space: None | Literal['Cuda', 'HIP', 'OpenMP'] | pykokkos.ExecutionSpace = None,

) → ModuleType

Return the Numpy or Cupy array module used for a specific execution space.

Parameters:

execution_space (None | {"Cuda", "HIP", "OpenMP"} | pk.ExecutionSpace) – Kokkos execution space specification. If None, return the Kokkos default execution space. The default is None.

Return type:

typing.ModuleType

Examples

>>> import parkipy
>>> am = parkipy.utils.get_array_module("Cuda")
>>> am
... <module 'cupy' from '/opt/homebrew/lib/python3.13/site-packages/cupy/__init__.py'>
>>> array = am.random.rand(3, 100000)