Incorporate GRIT model

[ttolhurst]

Incorporation of GRIT from L. Ma's "Graph Inductive Biases in Transformers without Message Passing".

[romeokienzler]

Code Review Findings

1. gridfm_graphkit/models/grit_transformer.py — GritHeteroAdapter.aggregate_pg reads self.grit.mask_value[0].item() to fill fully-masked buses. Since mask_value is registered as a learnable nn.Parameter when learn_mask=True, .item() detaches it: the learnable mask token receives no gradient through the only place it's consumed. learn_mask=True is silently a no-op.

2. gridfm_graphkit/training/loss.py — PBELoss computes S_injection = torch.diag(V) @ Y_bus_conj @ V_conj. torch.diag(V) materializes a dense N×N matrix and the matmul densifies the sparse Y-bus; on case2000 this is O(N²) memory and defeats the sparse Y-bus construction directly above. Use V * (Y_bus_conj @ V_conj) instead.

3. gridfm_graphkit/training/loss.py — PBELoss builds Y_diag = Y_diag + bus_orig[:, GS] + 1j * bus_orig[:, BS], pulling shunt admittance from the normalized input x_dict[\"bus\"]. If the normalizer scales GS/BS by a different factor than the Y-bus edge attributes YFF/YFT, the assembled Y-bus is physically inconsistent and the PBE residual is biased.

4. gridfm_graphkit/datasets/rrwp.py — Walk-length semantics are inconsistent with add_identity: with add_identity=True the PE has walk_length powers [I, A, A², …]; with add_identity=False it has walk_length-1 powers. Downstream RRWPLinearNodeEncoder(emb_dim=ksteps) assumes one fixed dimension, so toggling add_identity produces a silent size mismatch.

5. gridfm_graphkit/models/grit_layer.py — torch_scatter is imported under try/except and stored as None, but pyg_softmax/propagate_attention call scatter, scatter_max, scatter_add directly without a guard. A missing torch-scatter install raises TypeError: 'NoneType' is not callable instead of the explicit ImportError pattern used in rrwp_encoder._check_scatter.

6. gridfm_graphkit/training/loss.py — MaskedGenMSE was changed from pred_dict[\"gen\"][mask_dict[\"gen\"][:, :(PG_H+1)]] to slicing pred/target first then masking. This is a behavior change for any config with output_gen_dim > 1 (the old form would broadcast/error, the new form computes a different quantity). There's no test asserting either path; confirm intent and add coverage.

7. gridfm_graphkit/datasets/rrwp.py — deg = adj.sum(dim=1) followed by deg_inv = 1.0 / adj.sum(dim=1) recomputes the row sum; reuse deg. Minor but executes every PE computation (mitigated by caching).

8. gridfm_graphkit/models/rrwp_encoder.py — RRWPLinearEdgeEncoder.__init__ accepts fill_value and uses it to build padding, but then unconditionally sets self.fill_value = 0.0. The stored attribute (and __repr__) always reports 0.0 regardless of the argument; either honor it or drop the parameter.

9. Side note — GritHeteroAdapter.__init__ mutates the shared args object (args.model.input_dim, args.model.gt.dim_hidden, args.model.encoder.posenc_RWSE.kernel.times). Fine for a single-model run, but it's a footgun for any future code that constructs multiple models from the same config.

[romeokienzler]

@ttolhurst had a go with claude code, can you please have a look at the found issues and comment?

[ttolhurst]

Hi Romeo, Thanks for the review, it looks good. I will dive into the details this week, and will let you know how it goes. Do not hesitate if there is anything else in the meantime. Best regards, Thomas

…

On Fri, 26 Jun 2026 at 16:57, Romeo Kienzler ***@***.***> wrote: *romeokienzler* left a comment (gridfm/gridfm-graphkit#22) <#22 (comment)> @ttolhurst <https://github.com/ttolhurst> had a go with claude code, can you please have a look at the found issues and comment? — Reply to this email directly, view it on GitHub <#22?email_source=notifications&email_token=AXWAGW2SPYOM7QGAR337C3D5B3PTTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBRGMZTQNJYG432M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4813385877>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AXWAGW3BSXD7ZSA4OFLBNL35B3PTTAVCNFSNUABGKJSXA33TNF2G64TZHMYTAMBXGE2TSMBZGU5US43TOVSTWMZWGM2DOMRWGM3DBILWAI> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AXWAGW5LGQW6J4SOHMQ2YJD5B3PTTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBRGMZTQNJYG432M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AXWAGWZI2HHNKVXQNO73HUT5B3PTTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBRGMZTQNJYG432M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

[ttolhurst]

Hello @romeokienzler,

Thanks again for the review. All of the comments have been addressed in the new commits:

1. removed code related to learnable mask, which can now be implement from scratch in the future
2. implemented directly
3. not and issue for the current normalizer used
4. fixed to always add identity
5. enforce scatter install
6. code pertains to HGNS only
7. implemented directly
8. implemented directly
9. mitigated with deepcopy

Best, regards

[romeokienzler]

thanks @ttolhurst

• merged