run_concert_mouse_colon.ipynb` fails: `CONCERT.__init__()` got unexpected keyword `allow_batch_kernel_scale

[The-fourth-Marquis]

Thanks for making CONCERT available and for maintaining the notebooks—this repo has been really helpful for reproducing and extending the workflow. While running run_concert_mouse_colon.ipynb I ran into a small API mismatch that stops execution; sharing details below in case it’s an easy fix and to help keep the notebook in sync with the current CONCERT implementation.

Summary

Running notebooks/run_concert_mouse_colon.ipynb crashes when constructing the model because the notebook passes allow_batch_kernel_scale=..., but the CONCERT constructor does not accept that argument.

Where it happens

• Notebook: notebooks/run_concert_mouse_colon.ipynb
• Model import in notebook: from concert_batch_gut import CONCERT
• Implementation in repo: src/CONCERT-long/concert_batch_gut.py (class CONCERT(nn.Module))

Steps to reproduce

1. Open notebooks/run_concert_mouse_colon.ipynb
2. Run cells up through “## Setup module”
3. Execute the cell that builds the model:

allow_batch_kernel_scale = False
model = CONCERT(
    cell_atts=cell_atts,
    num_genes=adata.n_vars,
    input_dim=192,
    GP_dim=GP_dim,
    Normal_dim=Normal_dim,
    n_batch=n_batch,
    encoder_layers=encoder_layers,
    decoder_layers=decoder_layers,
    noise=noise,
    encoder_dropout=dropoutE,
    decoder_dropout=dropoutD,
    shared_dispersion=shared_dispersion,
    fixed_inducing_points=fix_inducing_points,
    initial_inducing_points=initial_inducing_points,
    fixed_gp_params=fixed_gp_params,
    kernel_scale=kernel_scale,
    allow_batch_kernel_scale=allow_batch_kernel_scale,  # <- triggers crash
    N_train=adata.n_obs,
    KL_loss=KL_loss,
    dynamicVAE=dynamicVAE,
    init_beta=init_beta,
    min_beta=min_beta,
    max_beta=max_beta,
    dtype=torch.float32,
    device=device,
)

Actual behavior

Notebook errors with:

TypeError: CONCERT.__init__() got an unexpected keyword argument 'allow_batch_kernel_scale'

Expected behavior

The notebook should run without crashing, either by:

• accepting allow_batch_kernel_scale (backwards compatible), or
• updating the notebook to use the current API.

Likely cause

src/CONCERT-long/concert_batch_gut.py defines CONCERT.__init__(..., kernel_scale, multi_kernel_mode, ...) but does not include allow_batch_kernel_scale. The notebook appears to be written against an older API (or a different CONCERT variant).

Proposed fix (pick one)

• Option A (preferred for UX/backwards compatibility): Add allow_batch_kernel_scale: bool = False to CONCERT.__init__ in src/CONCERT-long/concert_batch_gut.py and either:
  • ignore it (if no longer needed), or
  • use it to validate/reshape kernel_scale for batched kernels.
• Option B (simplest): Update notebooks/run_concert_mouse_colon.ipynb to remove allow_batch_kernel_scale=... and rely on the current behavior (since kernel_scale already supports float | np.ndarray in the current signature).
• Option C (docs): Document the expected kwargs for each CONCERT variant and align notebooks accordingly.

Environment

• OS: Linux (kernel 6.17.0-14-generic)
• Repo path: /root/code/26-02/CONCERT

python 3.11.11
torch 2.6.0+cu124
torch_cuda 12.4