Generate deterministic test tensors for all sizes and dtypes.
()
| 148 | # --------------------------------------------------------------------------- |
| 149 | |
| 150 | def generate_test_data() -> None: |
| 151 | """Generate deterministic test tensors for all sizes and dtypes.""" |
| 152 | |
| 153 | os.makedirs(TEST_DATA_DIR, exist_ok=True) |
| 154 | print("Generating test data...") |
| 155 | |
| 156 | gen = torch.Generator(device="cpu") |
| 157 | |
| 158 | for size_name, dims in MATMUL_SIZES: |
| 159 | M, N, K = dims["M"], dims["N"], dims["K"] |
| 160 | for dtype in TEST_DTYPES: |
| 161 | tag = _dtype_tag(dtype) |
| 162 | label = f" matmul/{size_name}/{tag}" |
| 163 | |
| 164 | save_dir = os.path.join(TEST_DATA_DIR, "matmul", size_name) |
| 165 | os.makedirs(save_dir, exist_ok=True) |
| 166 | save_path = os.path.join(save_dir, f"{tag}.pt") |
| 167 | |
| 168 | if os.path.exists(save_path): |
| 169 | print(f"{label} ... cached") |
| 170 | continue |
| 171 | |
| 172 | # Deterministic generation -- seed is fixed per (size, dtype) pair |
| 173 | gen.manual_seed(_SEED) |
| 174 | A = torch.randn(M, K, generator=gen, dtype=dtype) |
| 175 | B = torch.randn(K, N, generator=gen, dtype=dtype) |
| 176 | |
| 177 | torch.save({"A": A, "B": B}, save_path) |
| 178 | print(f"{label} ... ok") |
| 179 | |
| 180 | print() |
| 181 | |
| 182 | |
| 183 | # --------------------------------------------------------------------------- |