shared loading matrix W not be synchronized across samples

[maryjis]

Hi,

I am trying to understand the implementation of mNSF, especially how the shared gene loading matrix (W) is optimized across multiple samples.

From the model description, my understanding is that mNSF uses a shared loading matrix (W) across samples, while each sample has its own spatial latent GP factors. In other words, for sample (m),

$$ \Lambda_m = W \exp(F_m) $$

where (W) should be common to all samples.

However, when looking at the implementation, it seems that each sample-specific model object has its own separate TensorFlow variable self.W, initialized in pf_multiSample.py:

self.W = tf.Variable(
    rng.exponential(size=(J, L)),
    dtype=dtp,
    constraint=misc.make_nonneg,
    name="loadings"
)

Then, during multi-sample training in training_multiSample.py, I noticed code that appears to copy W from one sample model to another:

list_tro[ksample].model.W.assign(
    list_tro[ksample-1].model.W.numpy()
)

I also noticed that W_updated is computed as an average across samples:

W_updated = W_updated + (list_tro[ksample].model.W / nsample)

but I could not find where this averaged W_updated is assigned back to all sample models.

This makes me wonder whether (W) is actually being jointly optimized as a shared loading matrix across samples, or whether each sample model ends up maintaining its own sample-specific W, possibly only partially synchronized by copying.

Could you clarify whether this is intended behavior?

Specifically:

1. Should all sample models share exactly the same (W) during and after training?
2. Is W_updated supposed to be assigned back to each sample model after averaging?
3. If not, how is the shared-loading constraint enforced in the current implementation?
4. Is there a recommended way to verify after training that all sample models have the same learned (W)?

For example, I would expect something like this to return values close to zero if (W) is truly shared:

import numpy as np

Ws = [fit.W.numpy() for fit in list_fit_]

for i in range(len(Ws)):
    for j in range(i + 1, len(Ws)):
        print(i, j, np.max(np.abs(Ws[i] - Ws[j])))

This returned:

0 1 1.6429119

This suggests that the learned loading matrices are not identical across sample models after training.