ksuit.scripts.validate_pytorch_compatibilities¶

Classes¶

StabilityTestModel

A "kitchen sink" model designed to stress-test the exact kernels we use.

Functions¶

`run_cuda_assertions`()	Strictly checks if CUDA is available and configured for all libraries.
`generate_fake_data`(num_graphs, nodes_per_graph, ...)	Generates a batch of fake graph data directly on the GPU.
`set_deterministic_mode`(seed)	Sets PyTorch deterministic settings for reproducibility.
`main`()	Main stability test function.

Module Contents¶

ksuit.scripts.validate_pytorch_compatibilities.run_cuda_assertions()¶: Strictly checks if CUDA is available and configured for all libraries. Fails hard if any check fails.

class ksuit.scripts.validate_pytorch_compatibilities.StabilityTestModel(in_features, hidden_features, out_features, aggregation='mean')¶

Bases: torch.nn.Module

A “kitchen sink” model designed to stress-test the exact kernels we use. It uses: - radius_graph (torch_geometric.nn.pool) - knn_graph (torch_geometric.nn.pool) - scatter (torch_scatter - for message passing) - segment_csr (torch_scatter - for global pooling)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

aggregation = 'mean'¶

lin1¶

lin2¶

lin_out¶

radius = 1.0¶

k = 8¶

forward(x, pos, batch)¶

ksuit.scripts.validate_pytorch_compatibilities.generate_fake_data(num_graphs, nodes_per_graph, num_features, device)¶: Generates a batch of fake graph data directly on the GPU.

ksuit.scripts.validate_pytorch_compatibilities.set_deterministic_mode(seed)¶: Sets PyTorch deterministic settings for reproducibility.

ksuit.scripts.validate_pytorch_compatibilities.main()¶: Main stability test function.