<jupyter_start><jupyter_text>Prepare data batched by graph size<jupyter_code>node_feats_batched = dict()
adj_mats_batched = dict()
for k, v in node_feats.items():
    node_feats_batched[k] = np.stack(v)
    adj_mats_batched[k] = np.stack(adjacency_matrices[k])
def message_passing(feats, adjs):
    outputs = dict()
    for k in feats.keys():
        outputs[k] = batch_matmul(adjs[k], feats[k])
    return outputs

output = message_passing(node_feats_batched, adj_mats_batched)
output[19].shape
num_nodes = 13
node_feats_batched[num_nodes].shape, adj_mats_batched[num_nodes].shape<jupyter_output><empty_output><jupyter_text>Prepare data padded with zeros to largest graph size<jupyter_code>largest_graph_size = 20

for size, feats in node_feats.items():
    pad_size = largest_graph_size - size
    node_feats[size] = np.pad(
        feats, 
        pad_width=(
            # syntax is n_before, n_after
            (0, 0),  # sample dimension, do not touch
            (0, pad_size),  # node dimension, pad to largest size
            (0, 0),  # feats simension, do not touch
        )
    )
    
for size, adj in adjacency_matrices.items():
    pad_size = largest_graph_size - size
    adjacency_matrices[size] = np.pad(
        adj,
        pad_width=[
            # syntax is n_before, n_after
            (0, 0),  # sample dimension, do not touch
            (0, pad_size),  # node dimension, pad to largest size
            (0, pad_size),  # node dimension, pad to largest size
        ]
    )
node_feats[16].shape<jupyter_output><empty_output><jupyter_text>Now, stack them along the same axis.<jupyter_code>node_feats_stacked = []
for size, feats in node_feats.items():
    node_feats_stacked.append(feats)
node_feats_stacked = np.concatenate(node_feats_stacked, axis=0)

adj_mats_stacked = []
for size, adj in adjacency_matrices.items():
    adj_mats_stacked.append(adj)
adj_mats_stacked = np.concatenate(adj_mats_stacked, axis=0)
mp1 = batch_matmul(adj_mats_stacked, node_feats_stacked)
summation = np.sum(mp1, axis=1)
summation.shape  # n_samples by n_features
batch_matmul??<jupyter_output><empty_output><jupyter_text>Test EquivalenceMake sure that message passing with padding is equivalent to message passing without padding.<jupyter_code>F = np.array([[1, 0], [1, 1]])
A = np.array([[1, 0], [0, 1]])

M = np.dot(A, F)
M
pad_size = 2
F_pad = np.pad(
    F, 
    pad_width=[
        (0, pad_size),
        (0, 0),
    ]
)
F_pad

A_pad = np.pad(
    A,
    pad_width=[
        (0, pad_size),
        (0, pad_size),
    ]
)

A_pad
A_pad
M_pad = np.dot(A_pad, F_pad)
M_pad
# HERE ARE THE TESTS OF EQUIVALENCE!
assert np.all(np.pad(M, [(0, pad_size), (0, 0)]) == M_pad)
assert np.all(M_pad[:-pad_size, :] == M)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11786104037831939, 'fraction_numerical': 0.014186977082575483, 'mean_word_length': 3.6140939597315436, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '76910', 'license': 'MIT', 'n_tokens_mistral': 1073, 'n_tokens_neox': 1046, 'n_words': 231, 'path': 'examples/vectorized-graph-testing.ipynb', 'repo_name': 'ElArkk/fundl', 'size': '2749'}