keras/src/backend/jax/numpy.py/median
def median(x, axis=None, keepdims=False):
    # axis of jnp.median must be hashable
    if isinstance(axis, list):
        axis = tuple(axis)
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        x = cast(x, config.floatx())

    result = jnp.median(x, axis=axis, keepdims=keepdims)

    # TODO: with jax < 0.4.26 jnp.median failed to keepdims when axis is None
    if keepdims is True and axis is None:
        while result.ndim < x.ndim:
            result = jnp.expand_dims(result, axis=-1)
    return result