utils: internal utilities for developing

moraine.api.utils_ (source). These names are not re-exported: import them from moraine.api.utils_.

mjit

mjit(**options)

numba.jit with the compiled code cached in moraine's numba cache directory; use it (or ngjit, ngpjit) instead of cache=True

ngjit module-attribute

ngjit = mjit(nopython=True, nogil=True)

ngpjit module-attribute

ngpjit = mjit(nopython=True, nogil=True, parallel=True)

is_cuda_available

is_cuda_available()

get_n_cpus_avail

get_n_cpus_avail()

get_mem_avail

get_mem_avail()

Memory available to this process in bytes: the free memory of the machine, or less under a cgroup limit (e.g. a SLURM job).

get_array_module

get_array_module(array)

Examples

Run when this manual is built; the output follows each snippet.

import numpy as np
from moraine.api.utils_ import get_array_module, is_cuda_available, get_n_cpus_avail

print(is_cuda_available())                       # True only when CUDA_VISIBLE_DEVICES names a GPU and cupy imports
xp = get_array_module(np.zeros(3))               # numpy for numpy arrays, cupy for cupy arrays
print(xp.__name__)
print(get_n_cpus_avail() >= 1)

False
numpy
True

Functions dispatch on the kind of their input with get_array_module; mjit, ngjit and ngpjit are the numba decorators of moraine (compiled code cached by a hash of the sources, decision 0029).