Fix a partial feature coverage

[kmontemayor2-sc]

Summary

Lets a heterogeneous dataset with node/edge features on only a subset of its types train
successfully, instead of hanging on a swallowed KeyError inside the sampling coroutine.

Two independent changes:

1. Skip feature fetch for featureless types in the sampler collate (base_sampler.py).
BaseDistNeighborSampler._collate_fn previously fetched features for every sampled node/edge
type. For a heterogeneous dataset with features on only some types, a featureless-but-reachable
type raised a KeyError that GLT's ConcurrentEventLoop swallows (logging one line, no
traceback) — dropping every batch so the loader blocks forever. Both heterogeneous
feature-fetch loops now skip any type absent from DistFeature.local_feature (the per-type
feature dict). The use_all2all node path is left unguarded with a comment, since it is never
enabled here.

2. Normalize divergent cross-rank feature registration in the partitioner
(dist_partitioner.py, dist_range_partitioner.py). Direct DistPartitioner users may register
different feature/label/weight type sets per rank, which desynchronizes the per-type partitioning
collectives. New helpers all_gather each rank's {type: (trailing_shape, dtype)} schema and
per-type local entity counts (via graphlearn_torch.distributed.rpc.all_gather, since
torch.distributed is not generally initialized during partitioning), then per type in the
global union: backfill an empty (0, dim) tensor on a rank missing it with zero local rows;
raise if a rank is missing it but has nonzero local rows (unreconstructable); raise on
trailing-shape/dtype conflicts. Raise conditions are evaluated from the gathered data so every
rank raises symmetrically instead of one raising while the others hang. Normalization runs from
partition() and from the individual partition_node_features_and_labels() /
partition_edge_index_and_edge_features() methods, guarded by per-category flags to avoid
gathering twice. Also fixes partition_node_features_and_labels to partition the union of node
feature and label types, so a label-only node type is no longer dropped when another node type
carries features.

Tests

• Partial edge-/node-feature coverage regression tests for the standard and ABLP loaders, each
  bounded by a join-with-timeout since the pre-fix bug hangs rather than errors.
• Partitioner: count-conditional backfill, nonzero-rows and dim-conflict raises, node
  feature/label union, a zero-local-rows featured type, and a range-partitioner backfill.
• make type_check passes; the affected unit-test files pass
  (distributed_neighborloader_test.py, dist_ablp_neighborloader_test.py,
  distributed_partitioner_test.py).

Notes / out of scope

• A non-fatal partial-coverage warning at loader construction is separate, out-of-scope work and
  is not present on this branch (left as a follow-up).
• Building a DistDataset on a rank that receives zero edges of a type trips a pre-existing GLT
  empty-topology limitation (coo_to_csr calls max() on an empty tensor), and the range
  partition book cannot represent an empty per-rank edge range. Both are orthogonal to this
  change; the zero-local-rows test therefore asserts on the partition output's feature keys
  rather than a full dataset build.

[kmontemayor2-sc]

/all_test

[github-actions]

GiGL Automation

@ 15:52:30UTC : 🔄 Integration Test started.

@ 17:20:44UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 15:52:30UTC : 🔄 C++ Unit Test started.

@ 15:54:21UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 15:52:31UTC : 🔄 Python Unit Test started.

@ 17:02:38UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 15:52:31UTC : 🔄 Scala Unit Test started.

@ 16:01:05UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 15:52:33UTC : 🔄 Lint Test started.

@ 15:57:39UTC : ❌ Workflow failed.
Please check the logs for more details.

[github-actions]

GiGL Automation

@ 15:52:34UTC : 🔄 E2E Test started.

@ 17:19:43UTC : ❌ Workflow failed.
Please check the logs for more details.