fix: add the rest of the types

Author: beckermr

jax_galsim/core/utils.py

@@ -9,14 +9,31 @@
 import numpy as np
 from jax.tree_util import tree_flatten
 
-CONST_TYPES = (float, int, np.ndarray, np.int32, np.int64, np.float32, np.float64)
+CONST_TYPES = (
+    float,
+    int,
+    np.ndarray,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+    np.complex64,
+    np.complex128,
+)
 CONST_TYPES_WITH_JAX = CONST_TYPES + (
     jax.Array,
     jnp.array,
+    jnp.int8,
+    jnp.int16,
     jnp.int32,
     jnp.int64,
     jnp.float32,
     jnp.float64,
+    jnp.complex64,
+    jnp.complex128,
 )