keras/src/trainers/data_adapters/data_adapter_utils.py/tf_sparse_to_jax_sparse
def tf_sparse_to_jax_sparse(x):
    import jax.experimental.sparse as jax_sparse

    values = np.asarray(x.values)
    indices = np.asarray(x.indices)
    return jax_sparse.BCOO((values, indices), shape=x.shape)