Supporting multiple `layers`, `X`, `obsm` etc.

[ilan-gold]

Description of feature

This question came up in #90

Right now, the loader supports one label column and one "data" object. How can we support multiple "data" objects (see #49 for the label column part)?

I can see two routes:

1. Inspect the incoming AnnData object to the loader and load whatever is found within that is aligned along obs (i.e., any obsm, layers, X etc.)

ds = ZarrSparseDataset(
        batch_size=4096,
        chunk_size=32,
        preload_nchunks=512,
).add_anndata(my_anndata_with_layers_and_obsm)
next(iter(ds)) # return object contains both `obsm` and `layers`

2. Expose a RNG argument so two loaders can use RNGs seeded identically. This approach would force the user to create two objects, but would grossly simplify our internal code. I lean towards this because we can always create new APIs that compose two ZarrSparseDataset objects, for example, and handle this internally.

[giovp]

A question related to this is: is it possible to also yield var_names? or is it assumed that the var_names are all aligned acros chunks? and if they are, is there a schema validator for it?

[ilan-gold]

or is it assumed that the var_names are all aligned acros chunks? and if they are, is there a schema validator for it?

I opened #33 to start thinking about this problem - yes the assumption is that the var names are aligned but this 100% should be checked more thoroughly.

I think this dovetails with #98's goals a bit too - writing out some small metadata that indicates a folder of anndata's was preshuffled would give us confidence the var spaces are all aligned.

Could you highlight what the use of the var names would be? A little more preprocessing or data/munging before training?

[giovp]

var names would be? A little more preprocessing or data/munging before training?

var_names or var in general are useful for models with non-positional gene id information (e.g. transformer based that don't take the same gene at the same position all the time). I think related to #33 is also in general whether there is any schema validation to happen before training. For instance, what if the number of categories for a given obs column changes across chunks? that shouldn't happen, but I think it doesn't error atm.