Jax arrays have the attribute ndim instead of axis()

mlstm_kernels/jax/recurrent/native_sequence_scan.py contains a call to **jnp.array.axis()**, which does not exist. Replacing this instead with **jnp.array.ndim** fixed the error for me. Thanks.