"""Whisper base building blocks for recipes/whisper-base.py. Everything here is built from the pinned openai/whisper-base checkpoint loaded through transformers (WhisperForConditionalGeneration); the program wrappers reuse its nn.Modules (weights are shared, not copied) and only re-express the forward pass so that it lowers to one WebGPU delegate per program. Program contract (same split and I/O as the playground's previous whisper-base export): mel forward(frames[1,3000,400] f32) -> power_mel[1,3000,80] f32 frames: 30 s of 16 kHz audio, zero-padded to 480000 samples, reflect-padded by 200 on both ends, 3000 windows of 400 samples at hop 160 (= whisper.js buildFrames). power_mel = |rfft(hann_periodic * frame)|^2 @ mel_filters (no log; see host_log_mel). encoder forward(log_mel[1,80,3000] f32) -> encoder_hidden[1,1500,512] f32 (fp32 weights) cross_kv forward(encoder_hidden[1,1500,512]) -> 12 x [1,8,1500,64] f32, order k0,v0,...,k5,v5 (per decoder layer: encoder_attn.k_proj / v_proj, heads-first [B,H,S,D]) decoder_split forward(tok_embeds[1,64,512], k0, v0, ..., k5, v5) -> hidden[1,64,512] tok_embeds: token embeddings WITHOUT positions for positions 0..64 (zero rows beyond the live prefix); the program adds positions 0..63, runs causal self-attention over the 64 rows and returns final-LayerNorm hidden states (row p predicts token p+1). decode_stream forward(embed[1,1,512], input_pos[1] i64, k0, v0, ..., k5, v5) -> hidden[1,1,512] embed = token embedding + positional embedding of `input_pos` (added on the host); self-attention K/V cache [1,448,8,64] x 6 layers x {k,v} lives inside the program (mutable buffers owned by the delegate); feed positions 0,1,2,... in order. lmhead forward(hidden[1,1,512]) -> logits[1,1,51872] (tied to decoder.embed_tokens; rows 51865..51871 are padding with logit -1e9, see LMHeadProgram) Precision (whisper-base.py --precision): q4 (previous-export parity) cross_kv / decoder_split / decode_stream / lmhead linears are 4-bit groupwise weight-only (torchao IntxWeightOnlyConfig int4, group 64, symmetric, "affine" qparams -> et_vk.linear_q4gsw); int8 token-embedding table. mel and encoder are fp32. fp32 every program fp32 (aten.linear on 2-D views); fp32 token-embedding table. Host-side files (reproduced from the checkpoint by the recipe, consumed by the playground JS): whisper_embed.i8.bin int8 [51865,512] round(embed_tokens / scale), scale = rowwise absmax/127 (q4) whisper_embed.scale.f32.bin f32 [51865] the per-row scales (q4) whisper_embed.f32.bin f32 [51865,512] embed_tokens (fp32) whisper_pos_embed.f32.bin f32 [448,512] decoder.embed_positions.weight vocab.json upstream tokenizer vocab (byte-level BPE, detokenization only) """ import copy import io import numpy as np import torch import torch.nn.functional as F SAMPLE_RATE = 16000 N_FFT = 400 HOP = 160 FRAMES = 3000 N_SAMPLES = FRAMES * HOP PAD = N_FFT // 2 N_FREQ = N_FFT // 2 + 1 N_MELS = 80 ENC_LEN = 1500 D_MODEL = 512 N_HEADS = 8 HEAD_DIM = 64 N_LAYERS = 6 MAX_DEC = 64 # decoder_split rows; the playground's decode budget (prefix + generated) KV_LEN = 448 # decode_stream self-attention cache length (= max_target_positions) VOCAB = 51865 VOCAB_PADDED = 51872 # lmhead rows (multiple of 8, = the previous export) PAD_LOGIT = -1e9 TOK = { "sot": 50258, "eot": 50257, "en": 50259, "transcribe": 50359, "notimestamps": 50363, } PREFIX_EN = [TOK["sot"], TOK["en"], TOK["transcribe"], TOK["notimestamps"]] # -------------------------------------------------------------------------------------------- # Host-side math that the playground does in JS (kept here so validation runs the same steps). # -------------------------------------------------------------------------------------------- def build_frames(pcm): """Mirror of whisper.js buildFrames(): [3000, 400] float32 windows (no window function).""" pcm = np.asarray(pcm, dtype=np.float32) base = np.zeros(N_SAMPLES + 2 * PAD, dtype=np.float32) n = min(pcm.size, N_SAMPLES) base[PAD : PAD + n] = pcm[:n] for i in range(PAD): base[PAD - 1 - i] = base[PAD + 1 + i] base[PAD + N_SAMPLES + i] = base[PAD + N_SAMPLES - 2 - i] idx = np.arange(FRAMES)[:, None] * HOP + np.arange(N_FFT)[None, :] return base[idx] def host_log_mel(power_mel): """Host step between the mel and encoder programs (WhisperFeatureExtractor normalization). power_mel [1,3000,80] (or [3000,80]) -> log_mel [1,80,3000]: x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; transpose. float32 throughout (the JS equivalent uses Float32Array + Math.log10). """ p = np.asarray(power_mel, dtype=np.float32).reshape(FRAMES, N_MELS) x = np.log10(np.maximum(p, np.float32(1e-10))).astype(np.float32) x = np.maximum(x, x.max() - np.float32(8.0)) x = ((x + np.float32(4.0)) / np.float32(4.0)).astype(np.float32) return np.ascontiguousarray(x.T)[None] class HostTables: """Token-embedding lookup as the playground JS does it, plus the files that carry the tables. kind="int8": whisper_embed.i8.bin (int8 [51865,512], round(w / s)) + whisper_embed.scale.f32.bin (s = rowwise absmax / 127): byte-identical to the previous export's files. kind="fp32": whisper_embed.f32.bin (fp32 [51865,512], the checkpoint's embed_tokens). Both: whisper_pos_embed.f32.bin (fp32 [448,512], decoder.embed_positions).""" def __init__(self, model, kind): e = model.model.decoder.embed_tokens.weight.detach().float() self.kind = kind self.pos = model.model.decoder.embed_positions.weight.detach().float().numpy() if kind == "int8": scale = e.abs().amax(dim=1) / 127.0 self.table = torch.round(e / scale[:, None]).clamp(-127, 127).to(torch.int8).numpy() self.scale = scale.numpy().astype(np.float32) elif kind == "fp32": self.table = e.numpy().astype(np.float32) self.scale = None else: raise ValueError(kind) def embed(self, ids, position=None): e = self.table[ids].astype(np.float32) if self.scale is not None: e = e * self.scale[ids][..., None] if position is not None: e = e + self.pos[position] return e.astype(np.float32) def files(self): out = {} if self.kind == "int8": out["whisper_embed.i8.bin"] = self.table.tobytes() out["whisper_embed.scale.f32.bin"] = self.scale.astype(" [1, 8, s, 64].""" return x.reshape(1, s, N_HEADS, HEAD_DIM).transpose(1, 2) def _merge(o, s): """[1, 8, s, 64] -> [1, s, 512].""" return o.transpose(1, 2).reshape(1, s, D_MODEL) def attention(q, k, v, mask, impl): """q/k/v [1,H,S,D]; scale = 1/sqrt(64) = 0.125 (a power of two, so applying it to q as HF does or to q.k as the kernels do is bit-identical).""" if impl == "et_vk": import executorch.backends.vulkan.custom_ops_lib # noqa: F401 (torch.ops.et_vk.*) return torch.ops.et_vk.sdpa(q, k, v, mask, None) s = torch.matmul(q * 0.125, k.transpose(-2, -1)) if mask is not None: s = s + mask return torch.matmul(torch.softmax(s, dim=-1), v) class EncoderProgram(torch.nn.Module): """HF WhisperEncoder forward re-expressed with 2-D fp32 linears. attn="eager": matmul + softmax (lowered to aten.bmm/_softmax; activation memory is planned). attn="et_vk": the fused et_vk.sdpa op (what the previous export used); on the public WebGPU runtime each call allocates its own two [8,1500,1500] f32 scratch buffers (72 MB each). """ def __init__(self, model, attn="eager"): super().__init__() self.enc = model.model.encoder self.attn = attn def forward(self, log_mel): e = self.enc x = F.gelu(e.conv1(log_mel)) x = F.gelu(e.conv2(x)) # [1, 512, 1500] x = x.permute(0, 2, 1).reshape(ENC_LEN, D_MODEL) + e.embed_positions.weight for layer in e.layers: a = layer.self_attn h = layer.self_attn_layer_norm(x) q = _heads(a.q_proj(h), ENC_LEN) k = _heads(a.k_proj(h), ENC_LEN) v = _heads(a.v_proj(h), ENC_LEN) o = attention(q, k, v, None, "et_vk" if self.attn == "et_vk" else "torch") x = x + a.out_proj(o.transpose(1, 2).reshape(ENC_LEN, D_MODEL)) h = layer.final_layer_norm(x) x = x + layer.fc2(F.gelu(layer.fc1(h))) return e.layer_norm(x).reshape(1, ENC_LEN, D_MODEL) def _lin(m, x): """nn.Linear that lowers on the public WebGPU runtime: fp32 aten.linear only accepts 2-D inputs, so fp32 weights get a [-1, K] view; int4 (q4gsw) linears accept any rank and keep it.""" if type(m.weight.data) is torch.Tensor and x.dim() != 2: return m(x.reshape(-1, x.shape[-1])).reshape(*x.shape[:-1], m.out_features) return m(x) class CrossKVProgram(torch.nn.Module): def __init__(self, decoder): super().__init__() self.layers = decoder.layers def forward(self, enc): out = [] for layer in self.layers: a = layer.encoder_attn out.append(_heads(_lin(a.k_proj, enc), ENC_LEN)) out.append(_heads(_lin(a.v_proj, enc), ENC_LEN)) return tuple(out) def _mlp(layer, x): return x + _lin(layer.fc2, F.gelu(_lin(layer.fc1, layer.final_layer_norm(x)))) def _cross(layer, x, k, v, s, impl): a = layer.encoder_attn h = layer.encoder_attn_layer_norm(x) o = attention(_heads(_lin(a.q_proj, h), s), k, v, None, impl) return x + _lin(a.out_proj, _merge(o, s)) class DecoderSplitProgram(torch.nn.Module): """Full-recompute decoder over MAX_DEC rows with a causal mask (no KV cache).""" def __init__(self, decoder, attn="et_vk"): super().__init__() self.dec = decoder self.attn = attn self.register_buffer( "pos", decoder.embed_positions.weight.detach()[:MAX_DEC].clone() ) causal = torch.full((MAX_DEC, MAX_DEC), -1e30).triu(1) # same fill as the previous export # et_vk.sdpa on the WebGPU runtime requires the full [B, H, S_q, S_kv] mask. self.register_buffer( "mask", causal.expand(1, N_HEADS, MAX_DEC, MAX_DEC).contiguous() ) def forward(self, tok_embeds, *cross_kv): x = tok_embeds + self.pos for i, layer in enumerate(self.dec.layers): a = layer.self_attn h = layer.self_attn_layer_norm(x) q = _heads(_lin(a.q_proj, h), MAX_DEC) k = _heads(_lin(a.k_proj, h), MAX_DEC) v = _heads(_lin(a.v_proj, h), MAX_DEC) o = attention(q, k, v, self.mask, self.attn) x = x + _lin(a.out_proj, _merge(o, MAX_DEC)) x = _cross( layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], MAX_DEC, self.attn ) x = _mlp(layer, x) return self.dec.layer_norm(x) class DecodeStreamProgram(torch.nn.Module): """One token per call; self-attention K/V cache [1,448,8,64] per layer kept in the program. Self-attention uses llama.update_cache + llama.custom_sdpa (executorch.extension.llm.custom_ops), which the Vulkan partitioner fuses into sdpa_with_kv_cache with an et_vk.select_as_symint position. Cross-attention uses et_vk.sdpa (attn="et_vk") or matmul/softmax (attn="torch"). """ def __init__(self, decoder, attn="et_vk"): super().__init__() from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 self.dec = decoder self.attn = attn for i in range(N_LAYERS): self.register_buffer( f"k_cache_{i}", torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM), persistent=False, ) self.register_buffer( f"v_cache_{i}", torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM), persistent=False, ) def reset(self): for i in range(N_LAYERS): getattr(self, f"k_cache_{i}").zero_() getattr(self, f"v_cache_{i}").zero_() def forward(self, embed, input_pos, *cross_kv): pos = input_pos[0].item() torch._check_is_size(pos) torch._check(pos < KV_LEN) x = embed for i, layer in enumerate(self.dec.layers): a = layer.self_attn h = layer.self_attn_layer_norm(x) q = _lin(a.q_proj, h).view(1, 1, N_HEADS, HEAD_DIM) k = _lin(a.k_proj, h).view(1, 1, N_HEADS, HEAD_DIM) v = _lin(a.v_proj, h).view(1, 1, N_HEADS, HEAD_DIM) kc = getattr(self, f"k_cache_{i}") vc = getattr(self, f"v_cache_{i}") torch.ops.llama.update_cache(k, kc, pos) torch.ops.llama.update_cache(v, vc, pos) o = torch.ops.llama.custom_sdpa(q, kc, vc, pos, None, 0.0, True, None) x = x + _lin(a.out_proj, o.reshape(1, 1, D_MODEL)) x = _cross(layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], 1, self.attn) x = _mlp(layer, x) return self.dec.layer_norm(x) class LMHeadProgram(torch.nn.Module): """Tied output projection. The vocab is padded to `pad_to` rows: executorch 1.5.1's q4gsw lowering pads the scales to a multiple of 4 output channels and then fails its own shape check for N = 51865, so N must already be a multiple of 4 (the previous export used 51872). Padded rows have zero weights and a -1e9 bias, so they can never win an argmax.""" def __init__(self, model, pad_to=0): super().__init__() w = model.model.decoder.embed_tokens.weight.detach().clone() n = w.shape[0] b = torch.zeros(max(n, pad_to)) if pad_to and pad_to > n: w = torch.cat([w, w.new_zeros(pad_to - n, w.shape[1])]) b[n:] = PAD_LOGIT self.proj = torch.nn.Linear(D_MODEL, w.shape[0], bias=True) self.proj.weight = torch.nn.Parameter(w, requires_grad=False) self.proj.bias = torch.nn.Parameter(b, requires_grad=False) def forward(self, hidden): return _lin(self.proj, hidden) # -------------------------------------------------------------------------------------------- # Quantization. # -------------------------------------------------------------------------------------------- Q4_DESC = ( "torchao.quantization.quant_api.IntxWeightOnlyConfig(weight_dtype=torch.int4, " "granularity=PerGroup(64), intx_choose_qparams_algorithm='affine') applied with quantize_ to " "every nn.Linear of the program; symmetric, scale = group absmax / 7.5, q in [-8, 7]" ) def quantize_q4(module, group_size=64): from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ quantize_( module, IntxWeightOnlyConfig( weight_dtype=torch.int4, granularity=PerGroup(group_size), intx_choose_qparams_algorithm="affine", ), ) return module Q8_HEAD_DESC = ( "PT2E with executorch.backends.vulkan.quantizer.VulkanQuantizer: weights int8 per-output-channel " "symmetric (PerChannelMinMaxObserver, the quantizer's own weight spec), input activation int8 " "per-tensor affine STATIC (HistogramObserver) calibrated on lmhead inputs; lowers to " "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw" ) def q8_lmhead(lm, calib_hidden): """Static-activation int8 x per-channel int8 lmhead (the only 8-bit linear executorch 1.5.1's Vulkan lowering emits that the public WebGPU runtime registers). Returns the converted GraphModule and the (input scale, zero point, dequantized weight) needed for an equivalent portable module.""" from executorch.backends.vulkan.quantizer.vulkan_quantizer import ( get_symmetric_quantization_config, VulkanQuantizer, ) from torchao.quantization.pt2e import HistogramObserver from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from torchao.quantization.pt2e.quantizer import QuantizationConfig, QuantizationSpec base = get_symmetric_quantization_config(is_dynamic=False, weight_bits=8) act = QuantizationSpec( dtype=torch.int8, quant_min=-128, quant_max=127, qscheme=torch.per_tensor_affine, is_dynamic=False, observer_or_fake_quant_ctr=HistogramObserver, ) cfg = QuantizationConfig( input_activation=act, output_activation=None, weight=base.weight, bias=None, is_qat=False ) with torch.no_grad(): gm = torch.export.export(lm.eval(), (calib_hidden[:1],), strict=False).module() gm = prepare_pt2e(gm, VulkanQuantizer().set_global(cfg)) for i in range(calib_hidden.shape[0]): gm(calib_hidden[i : i + 1]) gm = convert_pt2e(gm) qp, wdq = None, None for n in gm.graph.nodes: if n.op != "call_function": continue name = str(n.target) if "quantize_per_tensor" in name and "dequantize" not in name and qp is None: qp = (float(n.args[1]), int(n.args[2])) if "dequantize_per_channel" in name: w = getattr(gm, n.args[0].target) sc = getattr(gm, n.args[1].target) wdq = (w.float() * sc.float()[:, None]).contiguous() if qp is None or wdq is None: raise RuntimeError("q8 lmhead: could not find the quantize/dequantize nodes") return gm, qp, wdq class FakeQ8LMHead(torch.nn.Module): """Portable-runtime equivalent of the q8 lmhead: per-tensor int8 fake-quant of the input, then an fp32 linear on the dequantized int8 weights (same bias).""" def __init__(self, lm, qp, wdq): super().__init__() self.scale, self.zp = qp self.register_buffer("w", wdq) self.register_buffer("b", lm.proj.bias.detach().clone()) def forward(self, hidden): x = hidden.reshape(-1, D_MODEL) q = torch.clamp(torch.round(x / self.scale) + self.zp, -128, 127) x = (q - self.zp) * self.scale return F.linear(x, self.w, self.b).reshape(*hidden.shape[:-1], self.w.shape[0]) def word_errors(tokenizer, hyp_text, ref_text): """(edits, reference words) after WhisperTokenizer.normalize (English text normalizer).""" h = tokenizer.normalize(hyp_text).split() r = tokenizer.normalize(ref_text).split() d = list(range(len(r) + 1)) for i in range(1, len(h) + 1): prev, d[0] = d[0], i for j in range(1, len(r) + 1): cur = d[j] d[j] = min(d[j] + 1, d[j - 1] + 1, prev + (h[i - 1] != r[j - 1])) prev = cur return d[len(r)], len(r) def dequantized_copy(module): """Same module with every quantized linear weight replaced by its fp32 dequantization. Used only for the CPU portable .pte (the portable runtime has no torchao kernels); its numerics are the q4 numerics (dequantize-then-fp32-matmul, as the q4gsw kernels compute).""" m = copy.deepcopy(module) for sub in m.modules(): if isinstance(sub, torch.nn.Linear) and hasattr(sub.weight, "dequantize"): sub.weight = torch.nn.Parameter( sub.weight.dequantize().float(), requires_grad=False ) return m # -------------------------------------------------------------------------------------------- # Speech sample. # -------------------------------------------------------------------------------------------- def load_librispeech_dummy(parquet_path, max_sec=15.0, count=8, longer=False): """First `count` clips (in file order) no longer than max_sec seconds, as 16 kHz float32. longer=True selects the clips *longer* than max_sec instead (used only for calibration, so calibration and evaluation clips never overlap).""" import pyarrow.parquet as pq import soundfile as sf table = pq.read_table(parquet_path).to_pylist() clips = [] for row in table: audio, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32") if sr != SAMPLE_RATE: raise ValueError(f"unexpected sample rate {sr}") if audio.ndim != 1: audio = audio.mean(axis=1) if (audio.size / SAMPLE_RATE > max_sec) != longer: continue clips.append({"id": row["id"], "text": row["text"], "pcm": audio}) if len(clips) >= count: break return clips def seconds(pcm): return round(len(pcm) / SAMPLE_RATE, 3) def isclose_report(a, b): a = np.asarray(a, dtype=np.float64).ravel() b = np.asarray(b, dtype=np.float64).ravel() d = np.abs(a - b) cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) return { "max_abs": float(d.max()), "mean_abs": float(d.mean()), "max_rel_to_absmax": float(d.max() / (np.abs(b).max() + 1e-30)), "cosine": cos, } # -------------------------------------------------------------------------------------------- # Lowered-graph inspection (what the WebGPU runtime will see). # -------------------------------------------------------------------------------------------- def _vk_graphs(pte_bytes): from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( extract_vk_flatbuffer, flatbuffer_to_vk_graph, ) from executorch.exir._serialize._program import deserialize_pte_binary pte = deserialize_pte_binary(pte_bytes) prog = pte.program if hasattr(pte, "program") else pte seg_base = int.from_bytes(pte_bytes[24:32], "little") out = [] for plan in prog.execution_plan: for d in plan.delegates: proc = d.processed if proc.location.name == "INLINE": blob = prog.backend_delegate_data[proc.index].data else: seg = prog.segments[proc.index] blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size] out.append(flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))) return prog, out def program_io(pte_bytes): """Method names and their input/output tensor shapes + dtypes from the serialized program.""" from executorch.exir._serialize._program import deserialize_pte_binary pte = deserialize_pte_binary(pte_bytes) prog = pte.program if hasattr(pte, "program") else pte res = {} for plan in prog.execution_plan: def desc(i): v = plan.values[i].val if hasattr(v, "sizes"): return {"shape": list(v.sizes), "dtype": v.scalar_type.name.lower()} return {"kind": type(v).__name__} res[plan.name] = { "inputs": [desc(i) for i in plan.inputs], "outputs": [desc(i) for i in plan.outputs], } return res def vk_checks(pte_bytes, runtime_ops=None): """Checks against constraints of the public WebGPU runtime's op implementations (backends/webgpu/runtime/ops): op registered, aten.linear 2-D, et_vk.sdpa mask shape, sdpa_with_kv_cache without mask / with causal, fp32 tensors only.""" _, graphs = _vk_graphs(pte_bytes) problems, notes = [], {} for g in graphs: vals = [v.value for v in g.values] def t(i): v = vals[i] return v if type(v).__name__ == "VkTensor" else None dtypes = sorted({type(v).__name__ if type(v).__name__ != "VkTensor" else v.datatype.name for v in vals}) notes["value_kinds"] = dtypes for op in g.chain: name = op.name if runtime_ops is not None and name not in runtime_ops: problems.append(f"{name}: not registered in the WebGPU runtime op list") a = op.args if name == "aten.linear.default": ri, rw = len(t(a[0]).dims), len(t(a[1]).dims) notes.setdefault("aten.linear input ranks", set()).add(ri) if ri != 2 or rw != 2: problems.append(f"aten.linear with input rank {ri}, weight rank {rw} (runtime needs 2-D)") if name == "et_vk.sdpa.default": q, k = t(a[0]), t(a[1]) m = t(a[3]) if m is not None: need = q.dims[0] * q.dims[1] * q.dims[2] * k.dims[2] got = int(np.prod(m.dims)) if need != got: problems.append(f"et_vk.sdpa mask numel {got} != B*H*Sq*Skv {need}") notes.setdefault("et_vk.sdpa q/k/mask", set()).add( (tuple(q.dims), tuple(k.dims), tuple(m.dims) if m is not None else None) ) if name == "sdpa_with_kv_cache.default": mask_kind = type(vals[a[7]]).__name__ causal = vals[a[9]] if mask_kind != "Null": problems.append("sdpa_with_kv_cache with attn_mask (runtime rejects)") if not (type(causal).__name__ == "Bool" and causal.bool_val): problems.append(f"sdpa_with_kv_cache is_causal={causal} (runtime needs true)") notes.setdefault("sdpa_with_kv_cache q/k_cache", set()).add( (tuple(t(a[0]).dims), tuple(t(a[3]).dims), type(vals[a[5]]).__name__) ) if name == "aten.convolution.default": notes.setdefault("aten.convolution input/weight", set()).add( (tuple(t(a[0]).dims), tuple(t(a[1]).dims)) ) for v in vals: if type(v).__name__ == "VkTensor" and v.datatype.name not in ("FLOAT32", "INT32", "INT64", "UINT8", "INT8"): problems.append(f"tensor dtype {v.datatype.name}") break for k, v in list(notes.items()): if isinstance(v, set): notes[k] = sorted(str(x) for x in v) return {"problems": problems, "ok": not problems, "notes": notes}