Bug: `ModuleNotFoundError: No module named 'torch_scatter'` when importing from `gridfm_graphkit`

Bug: ModuleNotFoundError: No module named 'torch_scatter' when importing from gridfm_graphkit

Summary

Any import from gridfm_graphkit — even from submodules that do not themselves use torch_scatter (e.g. gridfm_graphkit.datasets.postprocessing) — triggers a top-level ModuleNotFoundError for torch_scatter because the package __init__.py eagerly imports the entire task tree, which in turn imports torch_scatter unconditionally.

Reproduction

from gridfm_graphkit.datasets.postprocessing import (
    compute_branch_currents_kA,
    compute_loading,
)

Full Traceback

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../gridfm_graphkit/__init__.py", line 2, in <module>
    import gridfm_graphkit.tasks.base_task
  File ".../gridfm_graphkit/tasks/__init__.py", line 1, in <module>
    from gridfm_graphkit.tasks.pf_task import PowerFlowTask
  File ".../gridfm_graphkit/tasks/pf_task.py", line 21, in <module>
    from gridfm_graphkit.tasks.reconstruction_tasks import ReconstructionTask
  File ".../gridfm_graphkit/tasks/reconstruction_tasks.py", line 1, in <module>
    from gridfm_graphkit.io.param_handler import load_model, get_loss_function
  File ".../gridfm_graphkit/io/param_handler.py", line 2, in <module>
    from gridfm_graphkit.training.loss import MixedLoss
  File ".../gridfm_graphkit/training/__init__.py", line 1, in <module>
    from gridfm_graphkit.training.loss import (
  File ".../gridfm_graphkit/training/loss.py", line 6, in <module>
    from torch_scatter import scatter_add
ModuleNotFoundError: No module named 'torch_scatter'

Root Cause

torch_scatter is not listed as a dependency in pyproject.toml yet it is imported at module level in at least the following files:

• gridfm_graphkit/training/loss.py:6
• gridfm_graphkit/models/utils.py:3
• gridfm_graphkit/models/gnn_heterogeneous_gns.py:6
• gridfm_graphkit/tasks/se_task.py:22
• gridfm_graphkit/tasks/opf_task.py:40

Expected Behaviour

Either:

1. torch-scatter is added to [project.dependencies], or
2. the top-level init.py no longer eagerly imports the full task tree, so that lightweight submodules (e.g. datasets.postprocessing) can be used without pulling in heavy training/model dependencies.

Environment

gridfm-graphkit version: 0.8.1
Python: 3.10 – 3.12
Branch: chore/torch-scatter-dep

Start with pyproject.toml and gridfm_graphkit/init.py, then trace the imports named in the traceback through tasks/init.py, training/loss.py, and the listed model and task files. Reproduce the lightweight datasets.postprocessing import in an environment without torch_scatter, then verify that the chosen dependency or import strategy removes the unexpected failure.