keras/src/backend/jax/linalg.py/lu_factor
def lu_factor(x):
    lu_factor_fn = jsp.linalg.lu_factor
    if x.ndim > 2:
        for i in range(x.ndim - 2):
            lu_factor_fn = jax.vmap(lu_factor_fn)

    return lu_factor_fn(x)