src/rstor/utils.py

import numpy as np
try:
    import numba
except ImportError:
    numba = None
import torch

THREADS_PER_BLOCK = 32  # 32 or 16
DEFAULT_NUMPY_FLOAT_TYPE = np.float32
DEFAULT_CUDA_FLOAT_TYPE = numba.float32
DEFAULT_TORCH_FLOAT_TYPE = torch.float32


DEFAULT_NUMPY_INT_TYPE = np.int32
DEFAULT_CUDA_INT_TYPE = numba.int32