refactor: util routine

src/elisa/util.py → src/elisa/util/config.py

@@ -1,4 +1,4 @@
-"""Helper functions."""
+"""Helper functions for computation environment configuration."""
 from __future__ import annotations
 
 import warnings
@@ -10,7 +10,7 @@
 
 
 def jax_enable_x64(use_x64: bool) -> None:
-    """Changes the default float precision of array in JAX.
+    """Changes the default float precision of arrays in JAX.
 
     Parameters
     ----------

src/elisa/util/typing.py

@@ -0,0 +1,14 @@
+from typing import Union
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+FloatType = jnp.result_type(float)
+IntType = jnp.result_type(int)
+
+JAXFloat = Array
+PRNGKey = Array
+JAXArray = Array
+NumpyArray = np.ndarray
+Array = Union[NumpyArray, JAXArray]