It could be useful to additionally provide support for a sparse-representation of X using BCOO types in JAX. We might be able to use sparsify to get sparsity-handling for free on internal routines, and then provide an interface like susie_pca_sparse, or perhaps add a flag to susie_pca.
