From 90636818ab1ec0ef7698596caf167b48657a9632 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 8 Apr 2026 06:40:42 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20[Performance]=20Vectorize=20adjacen?= =?UTF-8?q?cy=20matrix=20construction=20in=20SymbolicDataset?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Vishal-sys-code <68536727+Vishal-sys-code@users.noreply.github.com> --- src/benchmarks/synthetic/symbolic_dataset.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/src/benchmarks/synthetic/symbolic_dataset.py b/src/benchmarks/synthetic/symbolic_dataset.py index aca5585..7367b5f 100644 --- a/src/benchmarks/synthetic/symbolic_dataset.py +++ b/src/benchmarks/synthetic/symbolic_dataset.py @@ -25,6 +25,7 @@ def __init__(self, num_samples: int = 2000, num_facts: int = 4): self.of = 7 self.entities = list(range(8, 34)) # A-Z + self.entities_t = torch.tensor(self.entities, dtype=torch.long) self.vocab_size = 35 self.data = [] @@ -65,19 +66,13 @@ def _generate_data(self): sequence.extend([self.query, e_a, self.is_t, self.anc, self.of, e_b]) - # Create Adjacency Matrix for "Same Entity" - seq_len = len(sequence) - A_rel = torch.zeros(seq_len, seq_len) - - # find all entities in sequence - for i in range(seq_len): - if sequence[i] in self.entities: - for j in range(seq_len): - if sequence[j] == sequence[i]: - A_rel[i, j] = 1.0 - x = torch.tensor(sequence, dtype=torch.long) y = torch.tensor(ans, dtype=torch.long) + + # Create Adjacency Matrix for "Same Entity" + mask = torch.isin(x, self.entities_t) + is_same = x.unsqueeze(1) == x.unsqueeze(0) + A_rel = (is_same & mask.unsqueeze(1)).float() self.data.append((x, y, A_rel))