Replace deprecated KeyArray with Array

[martinkim0]

JAX 0.4.24 removed random.KeyArray. According to the deprecation warning, it can be replaced with jax.Array:

DeprecationWarning: jax.random.KeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key) for runtime detection of typed prng keys (i.e. keys created with jax.random.key).

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (620f109) 91.00% compared to head (c5c0b0a) 91.01%.
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #135   +/-   ##
=======================================
  Coverage   91.00%   91.01%           
=======================================
  Files          25       25           
  Lines         956      957    +1     
=======================================
+ Hits          870      871    +1     
  Misses         86       86           

Files	Coverage Δ
src/scib_metrics/_types.py	100.00% <100.00%> (ø)
src/scib_metrics/utils/_kmeans.py	93.06% <100.00%> (+0.06%)	⬆️
src/scib_metrics/utils/_utils.py	60.60% <100.00%> (ø)