keras/src/backend/jax/numpy.py/concatenate
def concatenate(xs, axis=0):
    bcoo_count = builtins.sum(isinstance(x, jax_sparse.BCOO) for x in xs)
    if bcoo_count:
        if bcoo_count == len(xs):
            axis = canonicalize_axis(axis, len(xs[0].shape))
            return jax_sparse.bcoo_concatenate(xs, dimension=axis)
        else:
            xs = [
                x.todense() if isinstance(x, jax_sparse.JAXSparse) else x
                for x in xs
            ]
    return jnp.concatenate(xs, axis=axis)