Better typing

[manuelgloeckler]

Currently, we do not follow jax typing practices. Official suggestion are given here, which more or less suggest the following rules

import jax.numpy as jnp
from jax import Array
from jax.typing import ArrayLike
# Type inputs as ArrayLike (floats, np.array, jnp.array...)
# Type output as Array
def f(x: ArrayLike) -> Array:
     return jnp.pow(x,2)

Currently we often use jnp.ndarray which is an alias for Array.