diff --git a/spike/openxla/arch_execution_check.py b/spike/openxla/arch_execution_check.py new file mode 100644 index 000000000..6dce3b469 --- /dev/null +++ b/spike/openxla/arch_execution_check.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +"""Bounded token-exact execution check for the issue #499 dense arch pack. + +Proves that the mlxcel-xla Rust emitter's prefill + decode graphs, compiled with +IREE (llvm-cpu) and run on a real checkpoint's dequantized f32 weights, produce +the SAME greedy continuation as an HF fp32 oracle, WITHOUT building the Rust +`xla-iree` feature. It is the in-agent counterpart of `scripts/xla/validate_arch.sh` +(whose two execution gates need the native IREE build): the full token-exact / +serve gate is run post-merge; this proves architectural correctness cheaply on the +smallest checkpoint of a family. + +Method (weights fed identically to both sides, so the only variable is the graph): + 1. Dequantize the checkpoint to f32 and run the HF fp32 greedy oracle (eager + attention, matching the emitter's explicit softmax), recording the prompt's + last-position logits and the reference token ids. + 2. Emit the prefill + decode host-sampled logits graphs from the model's REAL + `config.json` via the scoped, pure-Rust `dump_dense_pack_graphs_for_execution_check` + test, and compile each with IREE (llvm-cpu). + 3. Load the f32 weights in the emitter's arg order (`weight_names.rs`), seed the + prefill graph, then drive the decode graph token by token, threading the KV + cache (mirroring `iree.rs::IreeLlama`), taking the argmax each step. + 4. Primary: the emitter's prefill logits must match HF's to a tight tolerance + (the whole forward on real weights). Secondary: the greedy token trajectory + (late divergences on an arbitrary prompt are near-tie flips, not errors). + +Only architectures the emitter supports load; a deferred arch is rejected at the +emit step (e.g. ERNIE-4.5, whose interleaved GPT-J RoPE this very check surfaced: +its prefill logits diverged from HF by ~8 while a supported family's match to +~1e-8). The in-agent supported checkpoints are large, so this is the opt-in / +post-merge counterpart to the cheap synthetic parity check (`synthetic_arch_check.py`); +it does NOT build the Rust `xla-iree` feature. + +Run (from the repo, using the spike venv's python), short prompt / CPU, chatty: + spike/openxla/.venv/bin/python spike/openxla/arch_execution_check.py \\ + --model /models/ --max-new 16 + +Exit 0 = forward parity PASS. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile + +import numpy as np + +import oracle_continuation as oc + +WORKTREE = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +MAX_SEQ = 256 # emitter MAX_SEQ / prefill bucket (PREFILL_LP) + + +def weight_arg_names(cfg: dict) -> list[str]: + """The emitter's weight arg order for the standard (Llama) scheme, mirroring + `weight_names.rs`: embed, final_norm, [lm_head if untied], then per layer + down, gate, in_ln, post_ln, up, wk, wo, wq, wv, [k/q/v bias if qkv_bias]. + + The dense pack's standard-name families (Seed-OSS, MiMo, InternLM3) load + through this; ExaOne 3.x's GPT-2-style scheme is a separate mapping not + exercised here. + """ + n_layers = int(cfg.get("num_hidden_layers") or cfg["num_layers"]) + tied = cfg.get("tie_word_embeddings", True) + mt = cfg.get("model_type") + qkv_bias = ( + mt == "qwen2" + or bool(cfg.get("attention_bias")) + or bool(cfg.get("qkv_bias")) + ) + names = ["model.embed_tokens.weight", "model.norm.weight"] + if not tied: + names.append("lm_head.weight") + for i in range(n_layers): + p = f"model.layers.{i}." + for suf in ( + "mlp.down_proj.weight", + "mlp.gate_proj.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + "mlp.up_proj.weight", + "self_attn.k_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_proj.weight", + "self_attn.v_proj.weight", + ): + names.append(p + suf) + if qkv_bias: + for suf in ( + "self_attn.k_proj.bias", + "self_attn.q_proj.bias", + "self_attn.v_proj.bias", + ): + names.append(p + suf) + return names + + +def hf_oracle( + weights_dir: str, prompt_ids: list[int], n_new: int +) -> tuple[list[int], np.ndarray]: + """HF fp32 greedy (pure argmax, no EOS stop) for `n_new` steps from explicit + `prompt_ids`, plus the prompt's last-position logits (the vector the emitter's + prefill returns). Returns `(ref_token_ids, prompt_last_logits)`. + + Eager attention is used so HF's softmax matches the emitter's explicit softmax + (the emitter has no fused SDPA), making the logit-parity comparison meaningful. + The tokenizer is deliberately not used: the check only needs the SAME prompt ids + on both sides, which sidesteps the custom tokenizers (and their optional deps) + several of these checkpoints ship. The caller strips the temp dequant config's + `auto_map` so HF loads the in-tree architecture class (no custom model code).""" + import torch + from transformers import AutoModelForCausalLM + + model = ( + AutoModelForCausalLM.from_pretrained( + weights_dir, dtype=torch.float32, attn_implementation="eager" + ) + .eval() + .float() + ) + ids = torch.tensor([prompt_ids], dtype=torch.long) + prompt_last_logits: np.ndarray | None = None + ref: list[int] = [] + with torch.no_grad(): + for step in range(n_new): + logits = model(ids).logits[:, -1, :] + if step == 0: + prompt_last_logits = logits[0].numpy().astype(np.float32) + nxt = int(torch.argmax(logits, dim=-1).item()) + ref.append(nxt) + ids = torch.cat([ids, torch.tensor([[nxt]])], dim=1) + assert prompt_last_logits is not None + return ref, prompt_last_logits + + +def emit_graphs(cfg_path: str, out_dir: str) -> tuple[str, str]: + """Dump the emitter's prefill + decode logits graphs from `cfg_path` via the + scoped, pure-Rust dump test. Returns (prefill_mlir, decode_mlir) paths.""" + print(f"[emit] cargo dump prefill/decode graphs from {cfg_path} ...", flush=True) + subprocess.run( + [ + "cargo", + "test", + "-p", + "mlxcel-xla", + "--lib", + "emitter::tests::dump_dense_pack_graphs_for_execution_check", + "--", + "--ignored", + "--nocapture", + ], + cwd=WORKTREE, + env={**os.environ, "MLXCEL_DUMP_CONFIG": cfg_path, "MLXCEL_DUMP_DIR": out_dir}, + check=True, + ) + return ( + os.path.join(out_dir, "prefill_logits.mlir"), + os.path.join(out_dir, "decode_logits.mlir"), + ) + + +def compile_vmfb(mlir_path: str, vmfb_path: str) -> None: + from iree.compiler.tools import compile_file + + print(f"[compile] iree-compile (llvm-cpu) {os.path.basename(mlir_path)} ...", flush=True) + compile_file( + mlir_path, + output_file=vmfb_path, + input_type="stablehlo", + target_backends=["llvm-cpu"], + ) + + +def load_f32_weights(dequant_dir: str, names: list[str]) -> list[np.ndarray]: + from safetensors import safe_open + + with safe_open(os.path.join(dequant_dir, "model.safetensors"), framework="numpy") as f: + avail = set(f.keys()) + out = [] + for n in names: + if n not in avail: + raise KeyError(f"weight {n} missing from dequantized checkpoint") + out.append(np.ascontiguousarray(f.get_tensor(n), dtype=np.float32)) + return out + + +def to_host(x) -> np.ndarray: + return x.to_host() if hasattr(x, "to_host") else np.asarray(x) + + +def emitter_continuation( + prefill_vmfb: str, + decode_vmfb: str, + weights: list[np.ndarray], + prompt_ids: list[int], + n_new: int, +) -> list[int]: + """Seed prefill then drive decode token by token (mirroring `iree.rs`).""" + from iree.runtime import load_vm_flatbuffer_file + + pre = load_vm_flatbuffer_file(prefill_vmfb, driver="local-task") + dec = load_vm_flatbuffer_file(decode_vmfb, driver="local-task") + + real_len = len(prompt_ids) + if real_len > MAX_SEQ: + raise ValueError(f"prompt of {real_len} exceeds MAX_SEQ={MAX_SEQ}") + tokens = np.zeros(MAX_SEQ, dtype=np.int32) + tokens[:real_len] = np.asarray(prompt_ids, dtype=np.int32) + positions = np.arange(MAX_SEQ, dtype=np.int32) + + print(f"[run] prefill (real_len={real_len}) ...", flush=True) + out = pre.main(*weights, tokens, positions, np.asarray(real_len, dtype=np.int32)) + logits = np.asarray(to_host(out[0]), dtype=np.float32) + prefill_logits = logits.copy() # the prompt's last-position logits + kc, vc = out[1], out[2] # keep as device arrays and thread them back in + tok = int(logits.argmax()) + gen = [tok] + + cache_len = real_len + for step in range(1, n_new): + out = dec.main( + *weights, + np.asarray(tok, dtype=np.int32), + np.asarray(cache_len, dtype=np.int32), + np.asarray(cache_len, dtype=np.int32), + kc, + vc, + ) + logits = to_host(out[0]) + kc, vc = out[1], out[2] + tok = int(np.asarray(logits).argmax()) + gen.append(tok) + cache_len += 1 + print(f"[run] decode step {step}: token={tok}", flush=True) + return gen, prefill_logits + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", required=True, help="checkpoint directory") + ap.add_argument( + "--prompt-ids", + default="", + help="comma-separated prompt token ids (default: a short valid sequence " + "derived from the config); the tokenizer is not used", + ) + ap.add_argument("--max-new", type=int, default=16) + args = ap.parse_args() + + with open(os.path.join(args.model, "config.json")) as f: + raw_cfg = json.load(f) + + work = tempfile.mkdtemp(prefix="arch_exec_") + # 1. Dequant (if needed) + HF fp32 oracle. + quant = oc.read_quantization(args.model) + dequant_dir = args.model + if quant is not None: + dequant_dir = os.path.join(work, "dequant") + print( + f"[oracle] dequantizing {args.model} (bits={quant['bits']}, " + f"group_size={quant['group_size']}) -> {dequant_dir}", + flush=True, + ) + oc.dequant_checkpoint(args.model, dequant_dir, quant["bits"], quant["group_size"]) + # Use the in-tree architecture class, not the checkpoint's bundled custom code: + # strip `auto_map` from the temp dequant config (safe only on the copy we made). + if dequant_dir != args.model: + dq_cfg_path = os.path.join(dequant_dir, "config.json") + with open(dq_cfg_path) as f: + dq_cfg = json.load(f) + if dq_cfg.pop("auto_map", None) is not None: + with open(dq_cfg_path, "w") as f: + json.dump(dq_cfg, f, indent=2) + # Prompt ids: explicit override, else a short valid sequence within vocab. + vocab = int(raw_cfg["vocab_size"]) + if args.prompt_ids.strip(): + prompt_ids = [int(x) for x in args.prompt_ids.split(",")] + else: + bos = raw_cfg.get("bos_token_id") + if isinstance(bos, list): + bos = bos[0] + seed = [bos, 1000, 2000, 4000, 8000, 16000] + prompt_ids = [t for t in seed if isinstance(t, int) and 0 <= t < vocab][:6] + if not prompt_ids or any(not (0 <= t < vocab) for t in prompt_ids): + raise ValueError(f"prompt ids {prompt_ids} out of range [0,{vocab})") + print(f"[oracle] prompt_ids={prompt_ids}", flush=True) + print("[oracle] HF fp32 greedy (eager) ...", flush=True) + ref, hf_prompt_logits = hf_oracle(dequant_dir, prompt_ids, args.max_new) + print(f"[oracle] ref continuation: {ref}", flush=True) + + # 2. Emit + compile the emitter graphs from the REAL config.json. + cfg_path = os.path.join(dequant_dir, "config.json") # dequant strips quantization + if not os.path.exists(cfg_path): + cfg_path = os.path.join(args.model, "config.json") + prefill_mlir, decode_mlir = emit_graphs(cfg_path, work) + prefill_vmfb = os.path.join(work, "prefill.vmfb") + decode_vmfb = os.path.join(work, "decode.vmfb") + compile_vmfb(prefill_mlir, prefill_vmfb) + compile_vmfb(decode_mlir, decode_vmfb) + + # 3. Load weights in emitter order, drive prefill + decode. + names = weight_arg_names(raw_cfg) + weights = load_f32_weights(dequant_dir, names) + print(f"[run] loaded {len(weights)} weight tensors in emitter arg order", flush=True) + got, emit_prompt_logits = emitter_continuation( + prefill_vmfb, decode_vmfb, weights, prompt_ids, len(ref) + ) + + # 4a. Primary criterion: prefill logit parity. The emitter's prefill runs the + # WHOLE forward (embed, every layer's attention + MLP + norms, RoPE, final norm, + # LM head) on the real dequantized weights; matching HF's last-position logits + # to a tight tolerance proves the architecture is emitted correctly. This is + # robust (unlike a bare argmax) to the ~1e-3 eager-vs-emitter numerical gap that + # can flip near-ties, especially on the arbitrary prompt ids used here. + diff = float(np.max(np.abs(emit_prompt_logits - hf_prompt_logits))) + rel = diff / float(np.max(np.abs(hf_prompt_logits)) + 1e-9) + ai, ah = int(emit_prompt_logits.argmax()), int(hf_prompt_logits.argmax()) + tol = 5e-2 + logit_ok = ai == ah and diff < tol + + # 4b. Secondary (informational): the greedy token trajectory. Reported as a + # match count; late divergences on the nonsensical prompt are near-tie flips, + # not an architecture error (the real token-exact gate uses real prompts and + # runs post-merge via the Rust IREE engine). + n = min(len(got), len(ref)) + traj_match = sum(1 for i in range(n) if got[i] == ref[i]) + print("", flush=True) + print(f"[compare] emitter tokens: {got[:n]}", flush=True) + print(f"[compare] oracle tokens: {ref[:n]}", flush=True) + print( + f"[compare] prefill logits: argmax emitter={ai} hf={ah} " + f"max|diff|={diff:.4e} (rel {rel:.2e}); trajectory {traj_match}/{n} match", + flush=True, + ) + if logit_ok: + print( + f"RESULT: PASS (prefill forward matches HF fp32: argmax equal, " + f"max|logit diff| {diff:.4e} < {tol})", + flush=True, + ) + return 0 + print( + f"RESULT: FAIL (prefill argmax emitter={ai} vs hf={ah}, " + f"max|logit diff| {diff:.4e})", + flush=True, + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/spike/openxla/synthetic_arch_check.py b/spike/openxla/synthetic_arch_check.py new file mode 100644 index 000000000..dd5dd3c3f --- /dev/null +++ b/spike/openxla/synthetic_arch_check.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +"""Synthetic HF-vs-emitter forward-parity check for the issue #499 dense pack. + +Some dense-pack families ship only in large or custom-code checkpoints that are +impractical to run in-agent, but have an in-tree HF config class. For those, this +builds a TINY random HF model (real modeling code, no download, no dequant), +freezes its weights, feeds them to the mlxcel-xla emitter's prefill graph compiled +with IREE (llvm-cpu), and compares last-token logits. The weights are identical on +both sides, so the only variable is the forward, and HF's real modeling code is the +oracle. This catches exactly the class of delta config inspection misses (e.g. +ERNIE-4.5's interleaved RoPE), which is why it complements the byte-exact structural +gate and the config/weight-name unit tests. + +Currently validates Seed-OSS (native `SeedOssConfig`; the untied + q/k/v-bias +Qwen2 forward with `rope_type = default`). It does NOT build the Rust `xla-iree` +feature (only the pure-Rust dump test), so it is watchdog-safe. + +Run (from the repo, using the spike venv's python): + spike/openxla/.venv/bin/python spike/openxla/synthetic_arch_check.py + +Exit 0 = PASS (argmax equal and last-token logits within tolerance). +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile + +import numpy as np +import torch +from iree.compiler.tools import compile_file +from iree.runtime import load_vm_flatbuffer_file + +WORKTREE = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +PREFILL_LP = 256 # emitter MAX_SEQ / prefill bucket +PROMPT_LEN = 24 +TOL = 2e-2 + + +def arg_order_names(n_layers: int, tied: bool, qkv_bias: bool) -> list[str]: + """Emitter weight arg order (Llama scheme) mirroring `weight_names.rs`.""" + names = ["model.embed_tokens.weight", "model.norm.weight"] + if not tied: + names.append("lm_head.weight") + for i in range(n_layers): + p = f"model.layers.{i}." + for suf in ( + "mlp.down_proj.weight", + "mlp.gate_proj.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + "mlp.up_proj.weight", + "self_attn.k_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_proj.weight", + "self_attn.v_proj.weight", + ): + names.append(p + suf) + if qkv_bias: + for suf in ( + "self_attn.k_proj.bias", + "self_attn.q_proj.bias", + "self_attn.v_proj.bias", + ): + names.append(p + suf) + return names + + +def emitter_last_logits(cfg_json: dict, weights_np: list[np.ndarray], prompt: np.ndarray) -> np.ndarray: + work = tempfile.mkdtemp(prefix="synth_arch_") + cfg_path = os.path.join(work, "config.json") + with open(cfg_path, "w") as fh: + json.dump(cfg_json, fh) + print("[emit] cargo dump prefill graph ...", flush=True) + subprocess.run( + [ + "cargo", "test", "-p", "mlxcel-xla", "--lib", + "emitter::tests::dump_dense_pack_graphs_for_execution_check", + "--", "--ignored", "--nocapture", + ], + cwd=WORKTREE, + env={**os.environ, "MLXCEL_DUMP_CONFIG": cfg_path, "MLXCEL_DUMP_DIR": work}, + check=True, + ) + vmfb = os.path.join(work, "prefill.vmfb") + print("[compile] iree-compile (llvm-cpu) ...", flush=True) + compile_file( + os.path.join(work, "prefill_logits.mlir"), + output_file=vmfb, + input_type="stablehlo", + target_backends=["llvm-cpu"], + ) + tokens = np.zeros(PREFILL_LP, dtype=np.int32) + tokens[:PROMPT_LEN] = prompt + positions = np.arange(PREFILL_LP, dtype=np.int32) + print("[run] IREE prefill ...", flush=True) + mod = load_vm_flatbuffer_file(vmfb, driver="local-task") + out = mod.main(*weights_np, tokens, positions, np.asarray(PROMPT_LEN, dtype=np.int32)) + logits = out[0].to_host() if hasattr(out[0], "to_host") else np.asarray(out[0]) + return np.asarray(logits, dtype=np.float32) + + +def check_seed_oss() -> bool: + from transformers import SeedOssConfig, SeedOssForCausalLM + + dims = dict( + hidden_size=16, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, # n_q*head_dim=32 != hidden=16 -> non-square o_proj, as in real Seed-OSS + intermediate_size=32, + num_hidden_layers=3, + vocab_size=64, + rms_norm_eps=1e-6, + rope_theta=10000.0, + max_position_embeddings=512, + attention_bias=True, + tie_word_embeddings=False, + ) + torch.manual_seed(0) + hf_cfg = SeedOssConfig(**dims) + hf_cfg._attn_implementation = "eager" + model = SeedOssForCausalLM(hf_cfg).eval().float() + with torch.no_grad(): + for _, p in model.named_parameters(): + if p.dim() == 1: # RMSNorm weights / biases: randomize so they are exercised + p.copy_(torch.randn_like(p) * 0.1) + state = {k: v.detach().clone() for k, v in model.state_dict().items()} + + names = arg_order_names(dims["num_hidden_layers"], tied=False, qkv_bias=True) + weights = [np.ascontiguousarray(state[n].numpy(), dtype=np.float32) for n in names] + + rng = np.random.default_rng(0) + prompt = rng.integers(0, dims["vocab_size"], size=PROMPT_LEN).astype(np.int32) + + emit_cfg = dict( + model_type="seed_oss", + hidden_size=dims["hidden_size"], + num_attention_heads=dims["num_attention_heads"], + num_key_value_heads=dims["num_key_value_heads"], + head_dim=dims["head_dim"], + intermediate_size=dims["intermediate_size"], + num_hidden_layers=dims["num_hidden_layers"], + vocab_size=dims["vocab_size"], + rms_norm_eps=dims["rms_norm_eps"], + rope_theta=dims["rope_theta"], + attention_bias=True, + tie_word_embeddings=False, + rope_scaling={"rope_type": "default"}, + ) + emit_logits = emitter_last_logits(emit_cfg, weights, prompt) + + with torch.no_grad(): + hf_logits = ( + model(input_ids=torch.tensor(prompt[None, :], dtype=torch.long)) + .logits[0, PROMPT_LEN - 1] + .numpy() + .astype(np.float32) + ) + diff = float(np.max(np.abs(emit_logits - hf_logits))) + ai, ah = int(emit_logits.argmax()), int(hf_logits.argmax()) + ok = ai == ah and diff < TOL + print( + f"[seed_oss] argmax emitter={ai} hf={ah} max|logit diff|={diff:.3e} -> " + f"{'OK' if ok else 'MISMATCH'}", + flush=True, + ) + return ok + + +def main() -> int: + ok = check_seed_oss() + print("RESULT:", "PASS" if ok else "FAIL", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lib/mlxcel-xla/assets/exaone/README.md b/src/lib/mlxcel-xla/assets/exaone/README.md new file mode 100644 index 000000000..85d74f123 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/exaone/README.md @@ -0,0 +1,23 @@ +# ExaOne 3.x config fixture (#499 dense arch pack) + +`config.json` is a **small synthetic** ExaOne config (real switches, tiny dims) +registered as a byte-exact structural fixture in `validation.rs`; the real +ExaOne-3.5 `config.json` parsing is covered by the `config::tests` parse assertion. +It deliberately uses the alternate field names ExaOne ships (`num_layers`, +`layer_norm_epsilon`, `activation_function`) so the fixture also exercises that +parsing. + +ExaOne 3.x (`model_type` `exaone`) is the proven **Llama** forward (llama3 RoPE, +tied embeddings, RMSNorm, SwiGLU) with a different tensor-naming scheme: the +GPT-2-style `transformer.h.{i}...` layout (the `Exaone` `WeightScheme`). Verified +from the checkpoint's `modeling_exaone.py`: half-split RoPE, `cat(freqs, freqs)` +table, `head_dim^-0.5` scaling, and a gated MLP `c_proj(act(c_fc_0(x)) * c_fc_1(x))` +(so `c_fc_0` is the gate, `c_fc_1` the up projection), attention under +`attn.attention.*` with `out_proj` as `o_proj`. + +The naming remap lives in `weight_names.rs` (unit-tested, +`exaone_scheme_maps_gpt2_style_names`); the emitted graph is loader-independent, so +`validation::dense_pack_families_reuse_proven_graphs` asserts ExaOne emits +byte-for-byte identical StableHLO to a tied llama3-RoPE Llama reference. The full +token-exact / serve gate runs post-merge on the real checkpoint via +`scripts/xla/validate_arch.sh`. diff --git a/src/lib/mlxcel-xla/assets/exaone/config.json b/src/lib/mlxcel-xla/assets/exaone/config.json new file mode 100644 index 000000000..c1b3c4e70 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/exaone/config.json @@ -0,0 +1,21 @@ +{ + "model_type": "exaone", + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 16, + "num_layers": 2, + "layer_norm_epsilon": 1e-05, + "rope_theta": 1000000.0, + "vocab_size": 10, + "tie_word_embeddings": true, + "activation_function": "silu", + "rope_scaling": { + "rope_type": "llama3", + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192 + } +} diff --git a/src/lib/mlxcel-xla/assets/exaone/decode.mlir b/src/lib/mlxcel-xla/assets/exaone/decode.mlir new file mode 100644 index 000000000..e5ab6e196 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/exaone/decode.mlir @@ -0,0 +1,212 @@ +module @decode_step { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg3: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg4: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg5: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg6: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg7: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg8: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg10: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg11: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg12: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg13: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg14: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg15: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg16: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg17: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg18: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg19: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg20: tensor loc("token"), %arg21: tensor loc("pos"), %arg22: tensor loc("cache_len"), %arg23: tensor<2x256x1x4xf32> loc("kcache"), %arg24: tensor<2x256x1x4xf32> loc("vcache")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3F0000803F40510A3F0000803F3311D5BEFEFF7F3F3311D5BEFEFF7F3F26707DBFFDFF7F3F26707DBFFDFF7F3F305527BFFAFF7F3F305527BFFAFF7F3F2C3C913EF6FF7F3F2C3C913EF6FF7F3FB8CD753FF2FF7F3FB8CD753FF2FF7F3FBDFF403FEDFF7F3FBDFF403FEDFF7F3FF6FD14BEE8FF7F3FF6FD14BEE8FF7F3FD53F69BFE1FF7F3FD53F69BFE1FF7F3F64CD56BFDAFF7F3F64CD56BFDAFF7F3F7205913BD2FF7F3F7205913BD2FF7F3FD006583FC9FF7F3FD006583FC9FF7F3F6F4E683FBFFF7F3F6F4E683FBFFF7F3FD7040C3EB5FF7F3FD7040C3EB5FF7F3FE87A42BFAAFF7F3FE87A42BFAAFF7F3F2C2975BF9EFF7F3F2C2975BF9EFF7F3F36E28CBE91FF7F3F36E28CBE91FF7F3F840A293F84FF7F3F840A293F84FF7F3FBF1B7D3F76FF7F3FBF1B7D3F76FF7F3F22F0D03E67FF7F3F22F0D03E67FF7F3FFC370CBF57FF7F3FFC370CBF57FF7F3F6FFD7FBF47FF7F3F6FFD7FBF47FF7F3FBF6708BF36FF7F3FBF6708BF36FF7F3FFE2DD93E24FF7F3FFE2DD93E24FF7F3F78BF7D3F11FF7F3F78BF7D3F11FF7F3F819C253FFDFE7F3F819C253FFDFE7F3F389395BEE9FE7F3F389395BEE9FE7F3F576D76BFD4FE7F3F576D76BFD4FE7F3FB3803FBFBEFE7F3FB3803FBFBEFE7F3F18F41D3EA8FE7F3F18F41D3EA8FE7F3F8E2C6A3F90FE7F3F8E2C6A3F90FE7F3FAA8F553F78FE7F3FAA8F553F78FE7F3FB78659BC5FFE7F3FB78659BC5FFE7F3FE73B59BF46FE7F3FE73B59BF46FE7F3F5F5867BF2BFE7F3F5F5867BF2BFE7F3FEA0803BE10FE7F3FEA0803BE10FE7F3F2DF2433FF4FD7F3F2DF2433FF4FD7F3FB57F743FD7FD7F3FB57F743FD7FD7F3F6C85883EBAFD7F3F6C85883EBAFD7F3F74BC2ABF9CFD7F3F74BC2ABF9CFD7F3F44C27CBF7DFD7F3F44C27CBF7DFD7F3FE0CACCBE5DFD7F3FE0CACCBE5DFD7F3FE81B0E3F3CFD7F3FE81B0E3F3CFD7F3FBBF57F3F1BFD7F3FBBF57F3F1BFD7F3F807B063FF9FC7F3F807B063FF9FC7F3F6D46DDBED6FC7F3F6D46DDBED6FC7F3FB3097EBFB2FC7F3FB3097EBFB2FC7F3F80E023BF8EFC7F3F80E023BF8EFC7F3F44E7993E69FC7F3F44E7993E69FC7F3F0308773F43FC7F3F0308773F43FC7F3FD1FD3D3F1CFC7F3FD1FD3D3F1CFC7F3F0EE726BEF5FB7F3F0EE726BEF5FB7F3F95146BBFCDFB7F3F95146BBFCDFB7F3FA64D54BFA4FB7F3FA64D54BFA4FB7F3F2C43B53C7AFB7F3F2C43B53C7AFB7F3FA26C5A3F50FB7F3FA26C5A3F50FB7F3FAC5D663F24FB7F3FAC5D663F24FB7F3FB714F43DF8FA7F3FB714F43DF8FA7F3F836545BFCCFA7F3F836545BFCCFA7F3F56D173BF9EFA7F3F56D173BF9EFA7F3FE62584BE70FA7F3FE62584BE70FA7F3FF76A2C3F41FA7F3FF76A2C3F41FA7F3FB8637C3F11FA7F3FB8637C3F11FA7F3F83A1C83EE0F97F3F83A1C83EE0F97F3FFAFC0FBFAFF97F3FFAFC0FBFAFF97F3FE5E87FBF7DF97F3FE5E87FBF7DF97F3F908C04BF4AF97F3F908C04BF4AF97F3F6D5AE13E16F97F3F6D5AE13E16F97F3FD54E7E3FE2F87F3FD54E7E3FE2F87F3F3521223FADF87F3F3521223FADF87F3F3A389EBE77F87F3F3A389EBE77F87F3FBC9D77BF40F87F3FBC9D77BF40F87F3F20773CBF08F87F3F20773CBF08F87F3FACD62F3ED0F77F3FACD62F3ED0F77F3FE4F76B3F97F77F3FE4F76B3F97F77F3F6107533F5DF77F3F6107533F5DF77F3F5ABFFDBC23F77F3F5ABFFDBC23F77F3FFC985BBFE7F67F3FFC985BBFE7F67F3F5A5E65BFABF67F3F5A5E65BFABF67F3FB512E2BD6EF67F3FB512E2BD6EF67F3FE4D4463F31F67F3FE4D4463F31F67F3F141E733FF2F57F3F141E733FF2F57F3F72877F3EB3F57F3F72877F3EB3F57F3F05162EBF73F57F3F05162EBF73F57F3F1C007CBF33F57F3F1C007CBF33F57F3F1F74C4BEF1F47F3F1F74C4BEF1F47F3F28DB113FAFF47F3F28DB113FAFF47F3FEDD67F3F6CF47F3FEDD67F3F6CF47F3FF69A023F28F47F3FF69A023F28F47F3FE869E5BEE4F37F3FE869E5BEE4F37F3FDE8E7EBF9EF37F3FDE8E7EBF9EF37F3FAA5E20BF58F37F3FAA5E20BF58F37F3F0486A23E12F37F3F0486A23E12F37F3F7D2E783FCAF27F3F7D2E783FCAF27F3FA8EC3A3F82F27F3FA8EC3A3F82F27F3FC2C238BE39F27F3FC2C238BE39F27F3F78D66CBFEFF17F3F78D66CBFEFF17F3FE0BC51BFA4F17F3FE0BC51BFA4F17F3F381B233D59F17F3F381B233D59F17F3FEEC05C3F0DF17F3FEEC05C3F0DF17F3F6E5A643FC0F07F3F6E5A643FC0F07F3F2A0CD03D72F07F3F2A0CD03D72F07F3F494048BF23F07F3F494048BF23F07F3FF16572BFD4EF7F3FF16572BFD4EF7F3FF8BD76BE84EF7F3FF8BD76BE84EF7F3F95BD2F3F34EF7F3F95BD2F3F34EF7F3F71977B3FE2EE7F3F71977B3FE2EE7F3FCB42C03E90EE7F3FCB42C03E90EE7F3F6AB613BF3DEE7F3F6AB613BF3DEE7F3FD4BF7FBFE9ED7F3FD4BF7FBFE9ED7F3FBFA600BF94ED7F3FBFA600BF94ED7F3FC874E93E3FED7F3FC874E93E3FED7F3FCBC97E3FE9EC7F3FCBC97E3FE9EC7F3FE7981E3F92EC7F3FE7981E3F92EC7F3F8BD0A6BE3AEC7F3F8BD0A6BE3AEC7F3F44BA78BFE2EB7F3F44BA78BFE2EB7F3F6F5E39BF89EB7F3F6F5E39BF89EB7F3F25AB413E2FEB7F3F25AB413E2FEB7F3F4CB06D3FD4EA7F3F4CB06D3FD4EA7F3F2A6E503F79EA7F3F2A6E503F79EA7F3F7E5347BD1CEA7F3F7E5347BD1CEA7F3F73E45DBFBFE97F3F73E45DBFBFE97F3FEE5163BF62E97F3FEE5163BF62E97F3F7301BEBD03E97F3F7301BEBD03E97F3FA9A7493FA4E87F3FA9A7493FA4E87F3FF2A8713F44E87F3FF2A8713F44E87F3F8BEF6D3EE3E77F3F8BEF6D3EE3E77F3F9F6131BF81E77F3F9F6131BF81E77F3FBB297BBF1FE77F3FBB297BBF1FE77F3F9C0DBCBEBCE67F3F9C0DBCBEBCE67F3FB68E153F58E67F3FB68E153F58E67F3F99A37F3FF3E57F3F99A37F3FF3E57F3FE45FFD3E8EE57F3FE45FFD3E8EE57F3FFA7AEDBE28E57F3FFA7AEDBE28E57F3F9DFF7EBFC1E47F3F9DFF7EBFC1E47F3FF6CF1CBF59E47F3FF6CF1CBF59E47F3FBA17AB3EF1E37F3FBA17AB3EF1E37F3F0E41793F87E37F3F0E41793F87E37F3F7FCC373F1DE37F3F7FCC373F1DE37F3FA48F4ABEB3E27F3FA48F4ABEB3E27F3F5B856EBF47E27F3F5B856EBF47E27F3F461B4FBFDBE17F3F461B4FBFDBE17F3FC5876B3D6EE17F3FC5876B3D6EE17F3F84035F3F00E17F3F84035F3F00E17F3FDF44623F91E07F3FDF44623F91E07F3FEDF2AB3D22E07F3FEDF2AB3D22E07F3FFD0A4BBFB2DF7F3FFD0A4BBFB2DF7F3F1AE770BF41DF7F3F1AE770BF41DF7F3F591C65BECFDE7F3F591C65BECFDE7F3F1A02333F5DDE7F3F1A02333F5DDE7F3FFCB67A3FEADD7F3FFCB67A3FEADD7F3FA7D4B73E76DD7F3FA7D4B73E76DD7F3F016417BF01DD7F3F016417BF01DD7F3F3E827FBF8BDC7F3F3E827FBF8BDC7F3F366DF9BE15DC7F3F366DF9BE15DC7F3F697CF13E9EDB7F3F697CF13E9EDB7F3F51307F3F26DB7F3F51307F3F26DB7F3FE0031B3FAEDA7F3FE0031B3FAEDA7F3F7A5BAFBE34DA7F3F7A5BAFBE34DA7F3FD8C279BFBAD97F3FD8C279BFBAD97F3FDF3636BF40D97F3FDF3636BF40D97F3F1470533EC4D87F3F1470533EC4D87F3FA2556F3F48D87F3FA2556F3F48D87F3F3BC44D3FCAD77F3F3BC44D3FCAD77F3FA9DB87BD4DD77F3FA9DB87BD4DD77F3F1C1E60BFCED67F3F1C1E60BFCED67F3F463361BF4ED67F3F463361BF4ED67F3FF3E099BDCED57F3FF3E099BDCED57F3F3F6A4C3F4DD57F3F3F6A4C3F4DD57F3F6D20703FCCD47F3F6D20703FCCD47F3F8E445C3E49D47F3F8E445C3E49D47F3FFE9E34BFC6D37F3FFE9E34BFC6D37F3F353F7ABF42D37F3F353F7ABF42D37F3F0298B3BEBDD27F3F0298B3BEBDD27F3F4336193F37D27F3F4336193F37D27F3FC35B7F3FB1D17F3FC35B7F3FB1D17F3F8875F53E2AD17F3F8875F53E2AD17F3FFF78F5BEA2D07F3FFF78F5BEA2D07F3FE75B7FBF1AD07F3FE75B7FBF1AD07F3FAE3419BF90CF7F3FAE3419BF90CF7F3FB69BB33E06CF7F3FB69BB33E06CF7F3F9F3F7A3F7BCE7F3F9F3F7A3F7BCE7F3F979D343FEFCD7F3F979D343FEFCD7F3F464C5CBE63CD7F3F464C5CBE63CD7F3F1C2170BFD6CC7F3F1C2170BFD6CC7F3F0F694CBF48CC7F3F0F694CBF48CC7F3FB6F0993DB9CB7F3FB6F0993DB9CB7F3F3634613F2ACB7F3F3634613F2ACB7F3F281D603F99CA7F3F281D603F99CA7F3FE4CB873D08CA7F3FE4CB873D08CA7F3F67C54DBF77C97F3F67C54DBF77C97F3FEF546FBFE4C87F3FEF546FBFE4C87F3F586853BE51C87F3F586853BE51C87F3F4238363FBDC77F3F4238363FBDC77F3F69C2793F28C77F3F69C2793F28C77F3FC357AF3E92C67F3FC357AF3E92C67F3F72051BBFFCC57F3F72051BBFFCC57F3F29307FBF65C57F3F29307FBF65C57F3FED78F1BECDC47F3FED78F1BECDC47F3FAA70F93E34C47F3FAA70F93E34C47F3F5D827F3F9BC37F3F5D827F3F9BC37F3F6962173F00C37F3F6962173F00C37F3F57D8B7BE65C27F3F57D8B7BE65C27F3F62B77ABFCAC17F3F62B77ABFCAC17F3FB00033BF2DC17F3FB00033BF2DC17F3F0D24653E90C07F3F0D24653E90C07F3FC5E7703FF2BF7F3FC5E7703FF2BF7F3FC9094B3F53BF7F3FC9094B3F53BF7F3FAC02ACBDB4BE7F3FAC02ACBDB4BE7F3FCB4562BF13BE7F3FCB4562BF13BE7F3F8C025FBF72BD7F3F8C025FBF72BD7F3F36686BBDD0BC7F3F36686BBDD0BC7F3F6F1C4F3F2EBC7F3F6F1C4F3F2EBC7F3FA4846E3F8ABB7F3FA4846E3F8ABB7F3FE5874A3EE6BA7F3FE5874A3EE6BA7F3FDFCD37BF41BA7F3FDFCD37BF41BA7F3F9A4079BF9CB97F3F9A4079BF9CB97F3F0014ABBEF5B87F3F0014ABBEF5B87F3F86D11C3F4EB87F3F86D11C3F4EB87F3F70FF7E3FA6B77F3F70FF7E3FA6B77F3F7A77ED3EFDB67F3F7A77ED3EFDB67F3F5363FDBE54B67F3F5363FDBE54B67F3FB4A37FBFA9B57F3FB4A37FBFA9B57F3F1B8D15BFFEB47F3F1B8D15BFFEB47F3F4911BC3E53B47F3F4911BC3E53B47F3F1D2A7B3FA6B37F3F1D2A7B3FA6B37F3F3260313FF9B27F3F3260313FF9B27F3F3BF76DBE4BB27F3F3BF76DBE4BB27F3F99A971BF9CB17F3F99A971BF9CB17F3F71A649BFECB07F3F71A649BFECB07F3F3011BE3D3CB07F3F3011BE3D3CB07F3FD752633F8BAF7F3FD752633F8BAF7F3F76E35D3FD9AE7F3F76E35D3FD9AE7F3FEC33473D26AE7F3FEC33473D26AE7F3F506F50BF73AD7F3F506F50BF73AD7F3F90AF6DBFBEAC7F3F90AF6DBFBEAC7F3F62A341BE0AAC7F3F62A341BE0AAC7F3FCC5F393F54AB7F3FCC5F393F54AB7F3FCCB9783F9DAA7F3FCCB9783F9DAA7F3FCFCCA63EE6A97F3FCFCCA63EE6A97F3F749A1EBF2EA97F3F749A1EBF2EA97F3F9AC97EBF75A87F3F9AC97EBF75A87F3F4471E9BEBCA77F3F4471E9BEBCA77F3F74A8003F01A77F3F74A8003F01A77F3FEABF7F3F46A67F3FEABF7F3F46A67F3FCDB4133F8AA57F3FCDB4133F8AA57F3F7546C0BECEA47F3F7546C0BECEA47F3FCF977BBF10A47F3FCF977BBF10A47F3F25BC2FBF52A37F3F25BC2FBF52A37F3FA3C5763E93A27F3FA3C5763E93A27F3F9466723FD4A17F3F9466723FD4A17F3F0E3F483F13A17F3F0E3F483F13A17F3FE31BD0BD52A07F3FE31BD0BD52A07F3F535B64BF909F7F3F535B64BF909F7F3FEEBF5CBFCD9E7F3FEEBF5CBFCD9E7F3F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573FD7FB5F39A46A573FD7FB5F39B7C7683FD6FBDF39B7C7683FD6FBDF39C381103EE0FC273AC381103EE0FC273ACFBD41BFD5FB5F3ACFBD41BFD5FB5F3A107C75BF64FD8B3A107C75BF64FD8B3A8C0F8FBEDEFCA73A8C0F8FBEDEFCA73A4630283F57FCC33A4630283F57FCC33A95467D3FD0FBDF3A95467D3FD0FBDF3A3201D33E47FBFB3A3201D33E47FBFB3AF8440BBF5FFD0B3BF8440BBF5FFD0B3B5CFF7FBF1AFD193B5CFF7FBF1AFD193BD85C09BFD5FC273BD85C09BFD5FC273B2220D73E8FFC353B2220D73E8FFC353B72987D3F49FC433B72987D3F49FC433B4479263F02FC513B4479263F02FC513B106893BEBAFB5F3B106893BEBAFB5F3B251E76BF72FB6D3B251E76BF72FB6D3BB34040BF29FB7B3BB34040BF29FB7B3B6979193E70FD843B6979193E70FD843BC8B6693F4AFD8B3BC8B6693F4AFD8B3B102F563F25FD923B102F563F25FD923B150511BCFEFC993B150511BCFEFC993BE7A158BFD8FCA03BE7A158BFD8FCA03BFCD367BFB1FCA73BFCD367BFB1FCA73B388707BE89FCAE3B388707BE89FCAE3B0837433F61FCB53B0837433F61FCB53B0ED5743F38FCBC3B0ED5743F38FCBC3B2AB48A3E0FFCC33B2AB48A3E0FFCC33BE9E329BFE5FBCA3BE9E329BFE5FBCA3BA4EF7CBFBBFBD13BA4EF7CBFBBFBD13B06DECEBE90FBD83B06DECEBE90FBD83B4D2A0D3F64FBDF3B4D2A0D3F64FBDF3B39FA7F3F38FBE63B39FA7F3F38FBE63BF771073F0BFBED3BF771073F0BFBED3BC23ADBBEDDFAF43BC23ADBBEDDFAF43B38E57DBFAFFAFB3B38E57DBFAFFAFB3BEABE24BF407D013CEABE24BF407D013C9FBD973E28FD043C9FBD973E28FD043C4BBB763F0F7D083C4BBB763F0F7D083CBCBF3E3FF7FC0B3CBCBF3E3FF7FC0B3CFB6D22BEDD7C0F3CFB6D22BEDD7C0F3C28A16ABFC4FC123C28A16ABFC4FC123C31EF54BFAA7C163C31EF54BFAA7C163CA103913C8FFC193CA103913C8FFC193CD0D4593F747C1D3CD0D4593F747C1D3C9ADB663F59FC203C9ADB663F59FC203CE813FD3D3D7C243CE813FD3D3D7C243C56AC44BF20FC273C56AC44BF20FC273C222974BF037C2B3C222974BF037C2B3CFF5586BEE6FB2E3CFF5586BEE6FB2E3C24942B3FC87B323C24942B3FC87B323CA0937C3FA9FB353CA0937C3FA9FB353CB4B6CA3E8A7B393CB4B6CA3E8A7B393CCD0C0FBF6AFB3C3CCD0C0FBF6AFB3C3CF4EF7FBF4A7B403CF4EF7FBF4A7B403C5E8405BF2AFB433C5E8405BF2AFB433CFD50DF3E087B473CFD50DF3E087B473CE72C7E3FE6FA4A3CE72C7E3FE6FA4A3C4301233FC47A4E3C4301233FC47A4E3C23109CBEA0FA513C23109CBEA0FA513C7F5377BF7D7A553C7F5377BF7D7A553CF23A3DBF58FA583CF23A3DBF58FA583C4B5F2B3E337A5C3C4B5F2B3E337A5C3CD4866B3F0DFA5F3CD4866B3F0DFA5F3C0BAB533FE779633C0BAB533FE779633CCE81D9BCC0F9663CCE81D9BCC0F9663C5B035BBF98796A3C5B035BBF98796A3C97DE65BF70F96D3C97DE65BF70F96D3C4D14EBBD4679713C4D14EBBD4679713CB31D463F1DF9743CB31D463F1DF9743C5178733FF278783C5178733FF278783C23F5813EC6F87B3C23F5813EC6F87B3CED402DBF9A787F3CED402DBF9A787F3C8C327CBF377C813C8C327CBF377C813C518BC6BE203C833C518BC6BE203C833C6EEC103F09FC843C6EEC103F09FC843C8EE07F3FF1BB863C8EE07F3FF1BB863C1894033FD97B883C1894033FD97B883CBC62E3BEC13B8A3CBC62E3BEC13B8A3C7D6F7EBFA8FB8B3C7D6F7EBFA8FB8B3C574021BF8EBB8D3C574021BF8EBB8D3C865FA03E757B8F3C865FA03E757B8F3CBCE6773F5B3B913CBCE6773F5B3B913C5CB23B3F40FB923C5CB23B3F40FB923C2B4D34BE25BB943C2B4D34BE25BB943CC6676CBF0A7B963CC6676CBF0A7B963CA76252BFEE3A983CA76252BFEE3A983CD0FD103DD1FA993CD0FD103DD1FA993C822D5C3FB4BA9B3C822D5C3FB4BA9B3CF7DC643F977A9D3CF7DC643F977A9D3CFB0FD93D793A9F3CFB0FD93D793A9F3C178B47BF5BFAA03C178B47BF5BFAA03C9EC272BF3CBAA23C9EC272BF3CBAA23C56237BBE1D7AA43C56237BBE1D7AA43C3DEA2E3FFE39A63C3DEA2E3FFE39A63C68CC7B3FDDF9A73C68CC7B3FDDF9A73CF25BC23EBDB9A93CF25BC23EBDB9A93C28C912BF9C79AB3C28C912BF9C79AB3C05CC7FBF7A39AD3C05CC7FBF7A39AD3C2EA101BF58F9AE3C2EA101BF58F9AE3CED6FE73E35B9B03CED6FE73E35B9B03CF8AC7E3F1279B23CF8AC7E3F1279B23C2F7C1F3FEE38B43C2F7C1F3FEE38B43CB1ABA4BEC9F8B53CB1ABA4BEC9F8B53C007578BFA5B8B73C007578BFA5B8B73C03263ABF7F78B93C03263ABF7F78B93C6D373D3E5938BB3C6D373D3E5938BB3CFA436D3F33F8BC3CFA436D3F33F8BC3C0B16513F0BB8BE3C0B16513F0BB8BE3CCF3735BDE477C03CCF3735BDE477C03C3E535DBFBB37C23C3E535DBFBB37C23CC0D663BF93F7C33CC0D663BF93F7C33C4F07C7BD69B7C53C4F07C7BD69B7C53C7AF4483F3F77C73C7AF4483F3F77C73C0D08723F1537C93C0D08723F1537C93C5D57723EE9F6CA3C5D57723EE9F6CA3C0B9030BFBDB6CC3C0B9030BFBDB6CC3C38617BBF9176CE3C38617BBF9176CE3CAE28BEBE6436D03CAE28BEBE6436D03CEFA2143F36F6D13CEFA2143F36F6D13C5BB27F3F08B6D33C5BB27F3F08B6D33C5557FF3ED975D53C5557FF3ED975D53C7878EBBEA935D73C7878EBBEA935D73C58E57EBF79F5D83C58E57EBF79F5D83CD4B41DBF48B5DA3CD4B41DBF48B5DA3C8FF4A83E1775DC3C8FF4A83E1775DC3C49FE783FE534DE3C49FE783FE534DE3CED95383FB2F4DF3CED95383FB2F4DF3CE41D46BE7EB4E13CE41D46BE7EB4E13C6C1B6EBF4A74E33C6C1B6EBF4A74E33C3DC54FBF1534E53C3DC54FBF1534E53C2D6E593DE0F3E63C2D6E593DE0F3E63C8A745E3FA9B3E83C8A745E3FA9B3E83CF8CB623F7273EA3CF8CB623F7273EA3CA4FAB43D3B33EC3CA4FAB43D3B33EC3CD5594ABF02F3ED3CD5594ABF02F3ED3CA14871BFC9B2EF3CA14871BFC9B2EF3C888669BE9072F13C888669BE9072F13C4F32323F5532F33C4F32323F5532F33CFCF07A3F1AF2F43CFCF07A3F1AF2F43C99F1B93EDEB1F63C99F1B93EDEB1F63CBC7916BFA171F83CBC7916BFA171F83C90937FBF6431FA3C90937FBF6431FA3C2F67FBBE26F1FB3C2F67FBBE26F1FB3C4B7CEF3EE7B0FD3C4B7CEF3EE7B0FD3C9B187F3FA770FF3C9B187F3FA770FF3C4FEA1B3F3398003D4FEA1B3F3398003D093AADBE1378013D093AADBE1378013D938279BFF257023D938279BFF257023D240237BFD037033D240237BFD037033D61004F3EAE17043D61004F3EAE17043D18EE6E3F8CF7043D18EE6E3F8CF7043D45704E3F6AD7053D45704E3F6AD7053D2EA07DBD47B7063D2EA07DBD47B7063D60915FBF2397073D60915FBF2397073DA3BC61BFFF76083DA3BC61BFFF76083D59EAA2BDDB56093D59EAA2BDDB56093D21BB4B3FB7360A3D21BB4B3FB7360A3D5D84703F92160B3D5D84703F92160B3D04B1603E6CF60B3D04B1603E6CF60B3DFFD033BF46D60C3DFFD033BF46D60C3DB97B7ABF20B60D3DB97B7ABF20B60D3DC9B6B5BEF9950E3DC9B6B5BEF9950E3D844D183FD2750F3D844D183FD2750F3DA46F7F3FAB55103DA46F7F3FAB55103DFE71F73E8335113DFE71F73E8335113D507BF3BE5A15123D507BF3BE5A15123DC0467FBF31F5123DC0467FBF31F5123DAA1C1ABF08D5133DAA1C1ABF08D5133D0A7CB13EDEB4143D0A7CB13EDEB4143DDC017A3FB494153DDC017A3FB494153DAF6A353F8A74163DAF6A353F8A74163DB7DE57BE5E54173DB7DE57BE5E54173DF9BB6FBF3334183DF9BB6FBF3334183D28174DBF0714193D28174DBF0714193D8CE6903DDAF3193D8CE6903DDAF3193DB9A9603FAED31A3DB9A9603FAED31A3DC7A8603F80B31B3DC7A8603F80B31B3DC9D6903D52931C3DC9D6903D52931C3D57184DBF24731D3D57184DBF24731D3D47BB6FBFF5521E3D47BB6FBFF5521E3DFED657BEC6321F3DFED657BEC6321F3D146C353F9612203D146C353F9612203D6F017A3F66F2203D6F017A3F66F2203D5578B13E35D2213D5578B13E35D2213D3E1E1ABF04B2223D3E1E1ABF04B2223D9A467FBFD391233D9A467FBFD391233DD777F3BEA071243DD777F3BEA071243D7375F73E6E51253D7375F73E6E51253DC66F7F3F3B31263DC66F7F3F3B31263DED4B183F0711273DED4B183F0711273D7BBAB5BED3F0273D7BBAB5BED3F0273D217C7ABF9ED0283D217C7ABF9ED0283D97CF33BF69B0293D97CF33BF69B0293DB9B8603E34902A3DB9B8603E34902A3D0B85703FFD6F2B3D0B85703FFD6F2B3DEFB94B3FC74F2C3DEFB94B3FC74F2C3D1AFAA2BD902F2D3D1AFAA2BD902F2D3D92BD61BF580F2E3D92BD61BF580F2E3D69905FBF20EF2E3D69905FBF20EF2E3DA2807DBDE7CE2F3DA2807DBDE7CE2F3D70714E3FAEAE303D70714E3FAEAE303D62ED6E3F748E313D62ED6E3F748E313DA4F84E3E396E323DA4F84E3E396E323D860337BFFF4D333D860337BFFF4D333D228279BFC32D343D228279BFC32D343D5136ADBE870D353D5136ADBE870D353DE0EB1B3F4BED353DE0EB1B3F4BED353D70187F3F0ECD363D70187F3F0ECD363DCD78EF3ED0AC373DCD78EF3ED0AC373DA06AFBBE928C383DA06AFBBE928C383DAD937FBF536C393DAD937FBF536C393D237816BF144C3A3D237816BF144C3A3D47F5B93ED42B3B3D47F5B93ED42B3B3D61F17A3F940B3C3D61F17A3F940B3C3DE430323F53EB3C3DE430323F53EB3C3D3A8E69BE12CB3D3D3A8E69BE12CB3D3D4A4971BFD0AA3E3D4A4971BFD0AA3E3D9F584ABF8D8A3F3D9F584ABF8D8A3F3D620AB53D4A6A403D620AB53D4A6A403DE3CC623F064A413DE3CC623F064A413D90735E3FC229423D90735E3FC229423D9D4E593D7D09433D9D4E593D7D09433D65C64FBF37E9433D65C64FBF37E9433DB31A6EBFF1C8443DB31A6EBFF1C8443D231646BEABA8453D231646BEABA8453D4C97383F6388463D4C97383F6388463DD3FD783F1C68473DD3FD783F1C68473DD4F0A83ED347483DD4F0A83ED347483D62B61DBF8A27493D62B61DBF8A27493D29E57EBF40074A3D29E57EBF40074A3DF674EBBEF6E64A3DF674EBBEF6E64A3DC15AFF3EABC64B3DC15AFF3EABC64B3D73B27F3F60A64C3D73B27F3F60A64C3D53A1143F14864D3D53A1143F14864D3D592CBEBEC7654E3D592CBEBEC7654E3D97617BBF7A454F3D97617BBF7A454F3D9D8E30BF2C25503D9D8E30BF2C25503D0B5F723EDD04513D0B5F723EDD04513DB208723F8EE4513DB208723F8EE4513D40F3483F3EC4523D40F3483F3EC4523D0917C7BDEEA3533D0917C7BDEEA3533DA7D763BF9D83543DA7D763BF9D83543D40525DBF4B63553D40525DBF4B63553D3C1835BDF942563D3C1835BDF942563D2F17513FA622573D2F17513FA622573D3C436D3F5202583D3C436D3F5202583DA92F3D3EFEE1583DA92F3D3EFEE1583D5E273ABFA9C1593D5E273ABFA9C1593D867478BF53A15A3D867478BF53A15A3DF3A7A4BEFD805B3DF3A7A4BEFD805B3DBA7D1F3FA6605C3DBA7D1F3FA6605C3DC5AC7E3F4F405D3DC5AC7E3F4F405D3D666CE73EF61F5E3D666CE73EF61F5E3DE2A201BF9EFF5E3DE2A201BF9EFF5E3D"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x3727C5AC> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.dynamic_slice %arg0, %arg20, %9, sizes = [1, 8] : (tensor<10x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %13 = stablehlo.reshape %12 : (tensor<1x8xf32>) -> tensor<8xf32> + %14 = stablehlo.dynamic_slice %0, %arg21, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %15 = stablehlo.reshape %14 : (tensor<1x4xf32>) -> tensor<4xf32> + %16 = stablehlo.dynamic_slice %1, %arg21, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %17 = stablehlo.reshape %16 : (tensor<1x4xf32>) -> tensor<4xf32> + %18 = stablehlo.iota dim = 0 : tensor<256xi32> + %19 = stablehlo.broadcast_in_dim %arg22, dims = [] : (tensor) -> tensor<256xi32> + %20 = stablehlo.compare LE, %18, %19, SIGNED : (tensor<256xi32>, tensor<256xi32>) -> tensor<256xi1> + %21 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256xf32> + %22 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256xf32> + %23 = stablehlo.select %20, %21, %22 : tensor<256xi1>, tensor<256xf32> + %24 = stablehlo.multiply %13, %13 : tensor<8xf32> + %25 = stablehlo.reduce(%24 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %26 = stablehlo.divide %25, %7 : tensor + %27 = stablehlo.add %26, %6 : tensor + %28 = stablehlo.rsqrt %27 : tensor + %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor) -> tensor<8xf32> + %30 = stablehlo.multiply %13, %29 : tensor<8xf32> + %31 = stablehlo.multiply %30, %arg4 : tensor<8xf32> + %32 = stablehlo.dot_general %31, %arg9, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %33 = stablehlo.reshape %32 : (tensor<8xf32>) -> tensor<2x4xf32> + %34 = stablehlo.dot_general %31, %arg7, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %35 = stablehlo.reshape %34 : (tensor<4xf32>) -> tensor<1x4xf32> + %36 = stablehlo.dot_general %31, %arg10, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %37 = stablehlo.reshape %36 : (tensor<4xf32>) -> tensor<1x4xf32> + %38 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %39 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %40 = stablehlo.multiply %33, %38 : tensor<2x4xf32> + %41 = stablehlo.slice %33 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %42 = stablehlo.slice %33 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %43 = stablehlo.negate %42 : tensor<2x2xf32> + %44 = stablehlo.concatenate %43, %41, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %45 = stablehlo.multiply %44, %39 : tensor<2x4xf32> + %46 = stablehlo.add %40, %45 : tensor<2x4xf32> + %47 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %48 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %49 = stablehlo.multiply %35, %47 : tensor<1x4xf32> + %50 = stablehlo.slice %35 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %51 = stablehlo.slice %35 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %52 = stablehlo.negate %51 : tensor<1x2xf32> + %53 = stablehlo.concatenate %52, %50, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %54 = stablehlo.multiply %53, %48 : tensor<1x4xf32> + %55 = stablehlo.add %49, %54 : tensor<1x4xf32> + %56 = stablehlo.reshape %55 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %57 = stablehlo.dynamic_update_slice %arg23, %56, %10, %arg22, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %58 = stablehlo.reshape %37 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %59 = stablehlo.dynamic_update_slice %arg24, %58, %10, %arg22, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %60 = stablehlo.slice %57 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %61 = stablehlo.reshape %60 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %62 = stablehlo.slice %59 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %63 = stablehlo.reshape %62 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %64 = stablehlo.reshape %46 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %65 = stablehlo.dot_general %64, %61, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %66 = stablehlo.reshape %65 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %67 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %68 = stablehlo.multiply %66, %67 : tensor<2x256xf32> + %69 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %70 = stablehlo.add %68, %69 : tensor<2x256xf32> + %71 = stablehlo.reduce(%70 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %72 = stablehlo.broadcast_in_dim %71, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %73 = stablehlo.subtract %70, %72 : tensor<2x256xf32> + %74 = stablehlo.exponential %73 : tensor<2x256xf32> + %75 = stablehlo.reduce(%74 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %76 = stablehlo.broadcast_in_dim %75, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %77 = stablehlo.divide %74, %76 : tensor<2x256xf32> + %78 = stablehlo.reshape %77 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %79 = stablehlo.dot_general %78, %63, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %80 = stablehlo.reshape %79 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %81 = stablehlo.reshape %80 : (tensor<2x4xf32>) -> tensor<8xf32> + %82 = stablehlo.dot_general %81, %arg8, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %83 = stablehlo.add %13, %82 : tensor<8xf32> + %84 = stablehlo.multiply %83, %83 : tensor<8xf32> + %85 = stablehlo.reduce(%84 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %86 = stablehlo.divide %85, %7 : tensor + %87 = stablehlo.add %86, %6 : tensor + %88 = stablehlo.rsqrt %87 : tensor + %89 = stablehlo.broadcast_in_dim %88, dims = [] : (tensor) -> tensor<8xf32> + %90 = stablehlo.multiply %83, %89 : tensor<8xf32> + %91 = stablehlo.multiply %90, %arg5 : tensor<8xf32> + %92 = stablehlo.dot_general %91, %arg3, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %93 = stablehlo.dot_general %91, %arg6, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %94 = stablehlo.negate %92 : tensor<16xf32> + %95 = stablehlo.exponential %94 : tensor<16xf32> + %96 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %97 = stablehlo.add %96, %95 : tensor<16xf32> + %98 = stablehlo.divide %96, %97 : tensor<16xf32> + %99 = stablehlo.multiply %92, %98 : tensor<16xf32> + %100 = stablehlo.multiply %99, %93 : tensor<16xf32> + %101 = stablehlo.dot_general %100, %arg2, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %102 = stablehlo.add %83, %101 : tensor<8xf32> + %103 = stablehlo.multiply %102, %102 : tensor<8xf32> + %104 = stablehlo.reduce(%103 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %105 = stablehlo.divide %104, %7 : tensor + %106 = stablehlo.add %105, %6 : tensor + %107 = stablehlo.rsqrt %106 : tensor + %108 = stablehlo.broadcast_in_dim %107, dims = [] : (tensor) -> tensor<8xf32> + %109 = stablehlo.multiply %102, %108 : tensor<8xf32> + %110 = stablehlo.multiply %109, %arg13 : tensor<8xf32> + %111 = stablehlo.dot_general %110, %arg18, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %112 = stablehlo.reshape %111 : (tensor<8xf32>) -> tensor<2x4xf32> + %113 = stablehlo.dot_general %110, %arg16, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %114 = stablehlo.reshape %113 : (tensor<4xf32>) -> tensor<1x4xf32> + %115 = stablehlo.dot_general %110, %arg19, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %116 = stablehlo.reshape %115 : (tensor<4xf32>) -> tensor<1x4xf32> + %117 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %118 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %119 = stablehlo.multiply %112, %117 : tensor<2x4xf32> + %120 = stablehlo.slice %112 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %121 = stablehlo.slice %112 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %122 = stablehlo.negate %121 : tensor<2x2xf32> + %123 = stablehlo.concatenate %122, %120, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %124 = stablehlo.multiply %123, %118 : tensor<2x4xf32> + %125 = stablehlo.add %119, %124 : tensor<2x4xf32> + %126 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %127 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %128 = stablehlo.multiply %114, %126 : tensor<1x4xf32> + %129 = stablehlo.slice %114 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %130 = stablehlo.slice %114 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %131 = stablehlo.negate %130 : tensor<1x2xf32> + %132 = stablehlo.concatenate %131, %129, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %133 = stablehlo.multiply %132, %127 : tensor<1x4xf32> + %134 = stablehlo.add %128, %133 : tensor<1x4xf32> + %135 = stablehlo.reshape %134 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %136 = stablehlo.dynamic_update_slice %57, %135, %11, %arg22, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %137 = stablehlo.reshape %116 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %138 = stablehlo.dynamic_update_slice %59, %137, %11, %arg22, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %139 = stablehlo.slice %136 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %140 = stablehlo.reshape %139 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %141 = stablehlo.slice %138 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %142 = stablehlo.reshape %141 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %143 = stablehlo.reshape %125 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %144 = stablehlo.dot_general %143, %140, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %145 = stablehlo.reshape %144 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %146 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %147 = stablehlo.multiply %145, %146 : tensor<2x256xf32> + %148 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %149 = stablehlo.add %147, %148 : tensor<2x256xf32> + %150 = stablehlo.reduce(%149 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %151 = stablehlo.broadcast_in_dim %150, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %152 = stablehlo.subtract %149, %151 : tensor<2x256xf32> + %153 = stablehlo.exponential %152 : tensor<2x256xf32> + %154 = stablehlo.reduce(%153 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %155 = stablehlo.broadcast_in_dim %154, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %156 = stablehlo.divide %153, %155 : tensor<2x256xf32> + %157 = stablehlo.reshape %156 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %158 = stablehlo.dot_general %157, %142, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %159 = stablehlo.reshape %158 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %160 = stablehlo.reshape %159 : (tensor<2x4xf32>) -> tensor<8xf32> + %161 = stablehlo.dot_general %160, %arg17, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %162 = stablehlo.add %102, %161 : tensor<8xf32> + %163 = stablehlo.multiply %162, %162 : tensor<8xf32> + %164 = stablehlo.reduce(%163 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %165 = stablehlo.divide %164, %7 : tensor + %166 = stablehlo.add %165, %6 : tensor + %167 = stablehlo.rsqrt %166 : tensor + %168 = stablehlo.broadcast_in_dim %167, dims = [] : (tensor) -> tensor<8xf32> + %169 = stablehlo.multiply %162, %168 : tensor<8xf32> + %170 = stablehlo.multiply %169, %arg14 : tensor<8xf32> + %171 = stablehlo.dot_general %170, %arg12, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %172 = stablehlo.dot_general %170, %arg15, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %173 = stablehlo.negate %171 : tensor<16xf32> + %174 = stablehlo.exponential %173 : tensor<16xf32> + %175 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %176 = stablehlo.add %175, %174 : tensor<16xf32> + %177 = stablehlo.divide %175, %176 : tensor<16xf32> + %178 = stablehlo.multiply %171, %177 : tensor<16xf32> + %179 = stablehlo.multiply %178, %172 : tensor<16xf32> + %180 = stablehlo.dot_general %179, %arg11, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %181 = stablehlo.add %162, %180 : tensor<8xf32> + %182 = stablehlo.multiply %181, %181 : tensor<8xf32> + %183 = stablehlo.reduce(%182 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %184 = stablehlo.divide %183, %7 : tensor + %185 = stablehlo.add %184, %6 : tensor + %186 = stablehlo.rsqrt %185 : tensor + %187 = stablehlo.broadcast_in_dim %186, dims = [] : (tensor) -> tensor<8xf32> + %188 = stablehlo.multiply %181, %187 : tensor<8xf32> + %189 = stablehlo.multiply %188, %arg1 : tensor<8xf32> + %190 = stablehlo.dot_general %189, %arg0, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %191 = stablehlo.constant dense<0> : tensor + %192 = stablehlo.constant dense<0xFF800000> : tensor + %193 = stablehlo.iota dim = 0 : tensor<10xi32> + %194:2 = stablehlo.reduce(%190 init: %192), (%193 init: %191) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %195 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %196 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %197 = stablehlo.or %195, %196 : tensor + %198 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %199 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %200 = stablehlo.and %198, %199 : tensor + %201 = stablehlo.or %197, %200 : tensor + %202 = stablehlo.select %197, %amv_l, %amv_r : tensor, tensor + %203 = stablehlo.select %201, %ami_l, %ami_r : tensor, tensor + stablehlo.return %202, %203 : tensor, tensor + } + return %194#1, %136, %138 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/exaone/prefill.mlir b/src/lib/mlxcel-xla/assets/exaone/prefill.mlir new file mode 100644 index 000000000..027173679 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/exaone/prefill.mlir @@ -0,0 +1,222 @@ +module @prefill { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg3: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg4: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg5: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg6: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg7: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg8: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg10: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg11: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg12: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg13: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg14: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg15: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg16: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg17: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg18: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg19: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg20: tensor<256xi32> loc("tokens"), %arg21: tensor<256xi32> loc("positions"), %arg22: tensor loc("real_len")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3F0000803F40510A3F0000803F3311D5BEFEFF7F3F3311D5BEFEFF7F3F26707DBFFDFF7F3F26707DBFFDFF7F3F305527BFFAFF7F3F305527BFFAFF7F3F2C3C913EF6FF7F3F2C3C913EF6FF7F3FB8CD753FF2FF7F3FB8CD753FF2FF7F3FBDFF403FEDFF7F3FBDFF403FEDFF7F3FF6FD14BEE8FF7F3FF6FD14BEE8FF7F3FD53F69BFE1FF7F3FD53F69BFE1FF7F3F64CD56BFDAFF7F3F64CD56BFDAFF7F3F7205913BD2FF7F3F7205913BD2FF7F3FD006583FC9FF7F3FD006583FC9FF7F3F6F4E683FBFFF7F3F6F4E683FBFFF7F3FD7040C3EB5FF7F3FD7040C3EB5FF7F3FE87A42BFAAFF7F3FE87A42BFAAFF7F3F2C2975BF9EFF7F3F2C2975BF9EFF7F3F36E28CBE91FF7F3F36E28CBE91FF7F3F840A293F84FF7F3F840A293F84FF7F3FBF1B7D3F76FF7F3FBF1B7D3F76FF7F3F22F0D03E67FF7F3F22F0D03E67FF7F3FFC370CBF57FF7F3FFC370CBF57FF7F3F6FFD7FBF47FF7F3F6FFD7FBF47FF7F3FBF6708BF36FF7F3FBF6708BF36FF7F3FFE2DD93E24FF7F3FFE2DD93E24FF7F3F78BF7D3F11FF7F3F78BF7D3F11FF7F3F819C253FFDFE7F3F819C253FFDFE7F3F389395BEE9FE7F3F389395BEE9FE7F3F576D76BFD4FE7F3F576D76BFD4FE7F3FB3803FBFBEFE7F3FB3803FBFBEFE7F3F18F41D3EA8FE7F3F18F41D3EA8FE7F3F8E2C6A3F90FE7F3F8E2C6A3F90FE7F3FAA8F553F78FE7F3FAA8F553F78FE7F3FB78659BC5FFE7F3FB78659BC5FFE7F3FE73B59BF46FE7F3FE73B59BF46FE7F3F5F5867BF2BFE7F3F5F5867BF2BFE7F3FEA0803BE10FE7F3FEA0803BE10FE7F3F2DF2433FF4FD7F3F2DF2433FF4FD7F3FB57F743FD7FD7F3FB57F743FD7FD7F3F6C85883EBAFD7F3F6C85883EBAFD7F3F74BC2ABF9CFD7F3F74BC2ABF9CFD7F3F44C27CBF7DFD7F3F44C27CBF7DFD7F3FE0CACCBE5DFD7F3FE0CACCBE5DFD7F3FE81B0E3F3CFD7F3FE81B0E3F3CFD7F3FBBF57F3F1BFD7F3FBBF57F3F1BFD7F3F807B063FF9FC7F3F807B063FF9FC7F3F6D46DDBED6FC7F3F6D46DDBED6FC7F3FB3097EBFB2FC7F3FB3097EBFB2FC7F3F80E023BF8EFC7F3F80E023BF8EFC7F3F44E7993E69FC7F3F44E7993E69FC7F3F0308773F43FC7F3F0308773F43FC7F3FD1FD3D3F1CFC7F3FD1FD3D3F1CFC7F3F0EE726BEF5FB7F3F0EE726BEF5FB7F3F95146BBFCDFB7F3F95146BBFCDFB7F3FA64D54BFA4FB7F3FA64D54BFA4FB7F3F2C43B53C7AFB7F3F2C43B53C7AFB7F3FA26C5A3F50FB7F3FA26C5A3F50FB7F3FAC5D663F24FB7F3FAC5D663F24FB7F3FB714F43DF8FA7F3FB714F43DF8FA7F3F836545BFCCFA7F3F836545BFCCFA7F3F56D173BF9EFA7F3F56D173BF9EFA7F3FE62584BE70FA7F3FE62584BE70FA7F3FF76A2C3F41FA7F3FF76A2C3F41FA7F3FB8637C3F11FA7F3FB8637C3F11FA7F3F83A1C83EE0F97F3F83A1C83EE0F97F3FFAFC0FBFAFF97F3FFAFC0FBFAFF97F3FE5E87FBF7DF97F3FE5E87FBF7DF97F3F908C04BF4AF97F3F908C04BF4AF97F3F6D5AE13E16F97F3F6D5AE13E16F97F3FD54E7E3FE2F87F3FD54E7E3FE2F87F3F3521223FADF87F3F3521223FADF87F3F3A389EBE77F87F3F3A389EBE77F87F3FBC9D77BF40F87F3FBC9D77BF40F87F3F20773CBF08F87F3F20773CBF08F87F3FACD62F3ED0F77F3FACD62F3ED0F77F3FE4F76B3F97F77F3FE4F76B3F97F77F3F6107533F5DF77F3F6107533F5DF77F3F5ABFFDBC23F77F3F5ABFFDBC23F77F3FFC985BBFE7F67F3FFC985BBFE7F67F3F5A5E65BFABF67F3F5A5E65BFABF67F3FB512E2BD6EF67F3FB512E2BD6EF67F3FE4D4463F31F67F3FE4D4463F31F67F3F141E733FF2F57F3F141E733FF2F57F3F72877F3EB3F57F3F72877F3EB3F57F3F05162EBF73F57F3F05162EBF73F57F3F1C007CBF33F57F3F1C007CBF33F57F3F1F74C4BEF1F47F3F1F74C4BEF1F47F3F28DB113FAFF47F3F28DB113FAFF47F3FEDD67F3F6CF47F3FEDD67F3F6CF47F3FF69A023F28F47F3FF69A023F28F47F3FE869E5BEE4F37F3FE869E5BEE4F37F3FDE8E7EBF9EF37F3FDE8E7EBF9EF37F3FAA5E20BF58F37F3FAA5E20BF58F37F3F0486A23E12F37F3F0486A23E12F37F3F7D2E783FCAF27F3F7D2E783FCAF27F3FA8EC3A3F82F27F3FA8EC3A3F82F27F3FC2C238BE39F27F3FC2C238BE39F27F3F78D66CBFEFF17F3F78D66CBFEFF17F3FE0BC51BFA4F17F3FE0BC51BFA4F17F3F381B233D59F17F3F381B233D59F17F3FEEC05C3F0DF17F3FEEC05C3F0DF17F3F6E5A643FC0F07F3F6E5A643FC0F07F3F2A0CD03D72F07F3F2A0CD03D72F07F3F494048BF23F07F3F494048BF23F07F3FF16572BFD4EF7F3FF16572BFD4EF7F3FF8BD76BE84EF7F3FF8BD76BE84EF7F3F95BD2F3F34EF7F3F95BD2F3F34EF7F3F71977B3FE2EE7F3F71977B3FE2EE7F3FCB42C03E90EE7F3FCB42C03E90EE7F3F6AB613BF3DEE7F3F6AB613BF3DEE7F3FD4BF7FBFE9ED7F3FD4BF7FBFE9ED7F3FBFA600BF94ED7F3FBFA600BF94ED7F3FC874E93E3FED7F3FC874E93E3FED7F3FCBC97E3FE9EC7F3FCBC97E3FE9EC7F3FE7981E3F92EC7F3FE7981E3F92EC7F3F8BD0A6BE3AEC7F3F8BD0A6BE3AEC7F3F44BA78BFE2EB7F3F44BA78BFE2EB7F3F6F5E39BF89EB7F3F6F5E39BF89EB7F3F25AB413E2FEB7F3F25AB413E2FEB7F3F4CB06D3FD4EA7F3F4CB06D3FD4EA7F3F2A6E503F79EA7F3F2A6E503F79EA7F3F7E5347BD1CEA7F3F7E5347BD1CEA7F3F73E45DBFBFE97F3F73E45DBFBFE97F3FEE5163BF62E97F3FEE5163BF62E97F3F7301BEBD03E97F3F7301BEBD03E97F3FA9A7493FA4E87F3FA9A7493FA4E87F3FF2A8713F44E87F3FF2A8713F44E87F3F8BEF6D3EE3E77F3F8BEF6D3EE3E77F3F9F6131BF81E77F3F9F6131BF81E77F3FBB297BBF1FE77F3FBB297BBF1FE77F3F9C0DBCBEBCE67F3F9C0DBCBEBCE67F3FB68E153F58E67F3FB68E153F58E67F3F99A37F3FF3E57F3F99A37F3FF3E57F3FE45FFD3E8EE57F3FE45FFD3E8EE57F3FFA7AEDBE28E57F3FFA7AEDBE28E57F3F9DFF7EBFC1E47F3F9DFF7EBFC1E47F3FF6CF1CBF59E47F3FF6CF1CBF59E47F3FBA17AB3EF1E37F3FBA17AB3EF1E37F3F0E41793F87E37F3F0E41793F87E37F3F7FCC373F1DE37F3F7FCC373F1DE37F3FA48F4ABEB3E27F3FA48F4ABEB3E27F3F5B856EBF47E27F3F5B856EBF47E27F3F461B4FBFDBE17F3F461B4FBFDBE17F3FC5876B3D6EE17F3FC5876B3D6EE17F3F84035F3F00E17F3F84035F3F00E17F3FDF44623F91E07F3FDF44623F91E07F3FEDF2AB3D22E07F3FEDF2AB3D22E07F3FFD0A4BBFB2DF7F3FFD0A4BBFB2DF7F3F1AE770BF41DF7F3F1AE770BF41DF7F3F591C65BECFDE7F3F591C65BECFDE7F3F1A02333F5DDE7F3F1A02333F5DDE7F3FFCB67A3FEADD7F3FFCB67A3FEADD7F3FA7D4B73E76DD7F3FA7D4B73E76DD7F3F016417BF01DD7F3F016417BF01DD7F3F3E827FBF8BDC7F3F3E827FBF8BDC7F3F366DF9BE15DC7F3F366DF9BE15DC7F3F697CF13E9EDB7F3F697CF13E9EDB7F3F51307F3F26DB7F3F51307F3F26DB7F3FE0031B3FAEDA7F3FE0031B3FAEDA7F3F7A5BAFBE34DA7F3F7A5BAFBE34DA7F3FD8C279BFBAD97F3FD8C279BFBAD97F3FDF3636BF40D97F3FDF3636BF40D97F3F1470533EC4D87F3F1470533EC4D87F3FA2556F3F48D87F3FA2556F3F48D87F3F3BC44D3FCAD77F3F3BC44D3FCAD77F3FA9DB87BD4DD77F3FA9DB87BD4DD77F3F1C1E60BFCED67F3F1C1E60BFCED67F3F463361BF4ED67F3F463361BF4ED67F3FF3E099BDCED57F3FF3E099BDCED57F3F3F6A4C3F4DD57F3F3F6A4C3F4DD57F3F6D20703FCCD47F3F6D20703FCCD47F3F8E445C3E49D47F3F8E445C3E49D47F3FFE9E34BFC6D37F3FFE9E34BFC6D37F3F353F7ABF42D37F3F353F7ABF42D37F3F0298B3BEBDD27F3F0298B3BEBDD27F3F4336193F37D27F3F4336193F37D27F3FC35B7F3FB1D17F3FC35B7F3FB1D17F3F8875F53E2AD17F3F8875F53E2AD17F3FFF78F5BEA2D07F3FFF78F5BEA2D07F3FE75B7FBF1AD07F3FE75B7FBF1AD07F3FAE3419BF90CF7F3FAE3419BF90CF7F3FB69BB33E06CF7F3FB69BB33E06CF7F3F9F3F7A3F7BCE7F3F9F3F7A3F7BCE7F3F979D343FEFCD7F3F979D343FEFCD7F3F464C5CBE63CD7F3F464C5CBE63CD7F3F1C2170BFD6CC7F3F1C2170BFD6CC7F3F0F694CBF48CC7F3F0F694CBF48CC7F3FB6F0993DB9CB7F3FB6F0993DB9CB7F3F3634613F2ACB7F3F3634613F2ACB7F3F281D603F99CA7F3F281D603F99CA7F3FE4CB873D08CA7F3FE4CB873D08CA7F3F67C54DBF77C97F3F67C54DBF77C97F3FEF546FBFE4C87F3FEF546FBFE4C87F3F586853BE51C87F3F586853BE51C87F3F4238363FBDC77F3F4238363FBDC77F3F69C2793F28C77F3F69C2793F28C77F3FC357AF3E92C67F3FC357AF3E92C67F3F72051BBFFCC57F3F72051BBFFCC57F3F29307FBF65C57F3F29307FBF65C57F3FED78F1BECDC47F3FED78F1BECDC47F3FAA70F93E34C47F3FAA70F93E34C47F3F5D827F3F9BC37F3F5D827F3F9BC37F3F6962173F00C37F3F6962173F00C37F3F57D8B7BE65C27F3F57D8B7BE65C27F3F62B77ABFCAC17F3F62B77ABFCAC17F3FB00033BF2DC17F3FB00033BF2DC17F3F0D24653E90C07F3F0D24653E90C07F3FC5E7703FF2BF7F3FC5E7703FF2BF7F3FC9094B3F53BF7F3FC9094B3F53BF7F3FAC02ACBDB4BE7F3FAC02ACBDB4BE7F3FCB4562BF13BE7F3FCB4562BF13BE7F3F8C025FBF72BD7F3F8C025FBF72BD7F3F36686BBDD0BC7F3F36686BBDD0BC7F3F6F1C4F3F2EBC7F3F6F1C4F3F2EBC7F3FA4846E3F8ABB7F3FA4846E3F8ABB7F3FE5874A3EE6BA7F3FE5874A3EE6BA7F3FDFCD37BF41BA7F3FDFCD37BF41BA7F3F9A4079BF9CB97F3F9A4079BF9CB97F3F0014ABBEF5B87F3F0014ABBEF5B87F3F86D11C3F4EB87F3F86D11C3F4EB87F3F70FF7E3FA6B77F3F70FF7E3FA6B77F3F7A77ED3EFDB67F3F7A77ED3EFDB67F3F5363FDBE54B67F3F5363FDBE54B67F3FB4A37FBFA9B57F3FB4A37FBFA9B57F3F1B8D15BFFEB47F3F1B8D15BFFEB47F3F4911BC3E53B47F3F4911BC3E53B47F3F1D2A7B3FA6B37F3F1D2A7B3FA6B37F3F3260313FF9B27F3F3260313FF9B27F3F3BF76DBE4BB27F3F3BF76DBE4BB27F3F99A971BF9CB17F3F99A971BF9CB17F3F71A649BFECB07F3F71A649BFECB07F3F3011BE3D3CB07F3F3011BE3D3CB07F3FD752633F8BAF7F3FD752633F8BAF7F3F76E35D3FD9AE7F3F76E35D3FD9AE7F3FEC33473D26AE7F3FEC33473D26AE7F3F506F50BF73AD7F3F506F50BF73AD7F3F90AF6DBFBEAC7F3F90AF6DBFBEAC7F3F62A341BE0AAC7F3F62A341BE0AAC7F3FCC5F393F54AB7F3FCC5F393F54AB7F3FCCB9783F9DAA7F3FCCB9783F9DAA7F3FCFCCA63EE6A97F3FCFCCA63EE6A97F3F749A1EBF2EA97F3F749A1EBF2EA97F3F9AC97EBF75A87F3F9AC97EBF75A87F3F4471E9BEBCA77F3F4471E9BEBCA77F3F74A8003F01A77F3F74A8003F01A77F3FEABF7F3F46A67F3FEABF7F3F46A67F3FCDB4133F8AA57F3FCDB4133F8AA57F3F7546C0BECEA47F3F7546C0BECEA47F3FCF977BBF10A47F3FCF977BBF10A47F3F25BC2FBF52A37F3F25BC2FBF52A37F3FA3C5763E93A27F3FA3C5763E93A27F3F9466723FD4A17F3F9466723FD4A17F3F0E3F483F13A17F3F0E3F483F13A17F3FE31BD0BD52A07F3FE31BD0BD52A07F3F535B64BF909F7F3F535B64BF909F7F3FEEBF5CBFCD9E7F3FEEBF5CBFCD9E7F3F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573FD7FB5F39A46A573FD7FB5F39B7C7683FD6FBDF39B7C7683FD6FBDF39C381103EE0FC273AC381103EE0FC273ACFBD41BFD5FB5F3ACFBD41BFD5FB5F3A107C75BF64FD8B3A107C75BF64FD8B3A8C0F8FBEDEFCA73A8C0F8FBEDEFCA73A4630283F57FCC33A4630283F57FCC33A95467D3FD0FBDF3A95467D3FD0FBDF3A3201D33E47FBFB3A3201D33E47FBFB3AF8440BBF5FFD0B3BF8440BBF5FFD0B3B5CFF7FBF1AFD193B5CFF7FBF1AFD193BD85C09BFD5FC273BD85C09BFD5FC273B2220D73E8FFC353B2220D73E8FFC353B72987D3F49FC433B72987D3F49FC433B4479263F02FC513B4479263F02FC513B106893BEBAFB5F3B106893BEBAFB5F3B251E76BF72FB6D3B251E76BF72FB6D3BB34040BF29FB7B3BB34040BF29FB7B3B6979193E70FD843B6979193E70FD843BC8B6693F4AFD8B3BC8B6693F4AFD8B3B102F563F25FD923B102F563F25FD923B150511BCFEFC993B150511BCFEFC993BE7A158BFD8FCA03BE7A158BFD8FCA03BFCD367BFB1FCA73BFCD367BFB1FCA73B388707BE89FCAE3B388707BE89FCAE3B0837433F61FCB53B0837433F61FCB53B0ED5743F38FCBC3B0ED5743F38FCBC3B2AB48A3E0FFCC33B2AB48A3E0FFCC33BE9E329BFE5FBCA3BE9E329BFE5FBCA3BA4EF7CBFBBFBD13BA4EF7CBFBBFBD13B06DECEBE90FBD83B06DECEBE90FBD83B4D2A0D3F64FBDF3B4D2A0D3F64FBDF3B39FA7F3F38FBE63B39FA7F3F38FBE63BF771073F0BFBED3BF771073F0BFBED3BC23ADBBEDDFAF43BC23ADBBEDDFAF43B38E57DBFAFFAFB3B38E57DBFAFFAFB3BEABE24BF407D013CEABE24BF407D013C9FBD973E28FD043C9FBD973E28FD043C4BBB763F0F7D083C4BBB763F0F7D083CBCBF3E3FF7FC0B3CBCBF3E3FF7FC0B3CFB6D22BEDD7C0F3CFB6D22BEDD7C0F3C28A16ABFC4FC123C28A16ABFC4FC123C31EF54BFAA7C163C31EF54BFAA7C163CA103913C8FFC193CA103913C8FFC193CD0D4593F747C1D3CD0D4593F747C1D3C9ADB663F59FC203C9ADB663F59FC203CE813FD3D3D7C243CE813FD3D3D7C243C56AC44BF20FC273C56AC44BF20FC273C222974BF037C2B3C222974BF037C2B3CFF5586BEE6FB2E3CFF5586BEE6FB2E3C24942B3FC87B323C24942B3FC87B323CA0937C3FA9FB353CA0937C3FA9FB353CB4B6CA3E8A7B393CB4B6CA3E8A7B393CCD0C0FBF6AFB3C3CCD0C0FBF6AFB3C3CF4EF7FBF4A7B403CF4EF7FBF4A7B403C5E8405BF2AFB433C5E8405BF2AFB433CFD50DF3E087B473CFD50DF3E087B473CE72C7E3FE6FA4A3CE72C7E3FE6FA4A3C4301233FC47A4E3C4301233FC47A4E3C23109CBEA0FA513C23109CBEA0FA513C7F5377BF7D7A553C7F5377BF7D7A553CF23A3DBF58FA583CF23A3DBF58FA583C4B5F2B3E337A5C3C4B5F2B3E337A5C3CD4866B3F0DFA5F3CD4866B3F0DFA5F3C0BAB533FE779633C0BAB533FE779633CCE81D9BCC0F9663CCE81D9BCC0F9663C5B035BBF98796A3C5B035BBF98796A3C97DE65BF70F96D3C97DE65BF70F96D3C4D14EBBD4679713C4D14EBBD4679713CB31D463F1DF9743CB31D463F1DF9743C5178733FF278783C5178733FF278783C23F5813EC6F87B3C23F5813EC6F87B3CED402DBF9A787F3CED402DBF9A787F3C8C327CBF377C813C8C327CBF377C813C518BC6BE203C833C518BC6BE203C833C6EEC103F09FC843C6EEC103F09FC843C8EE07F3FF1BB863C8EE07F3FF1BB863C1894033FD97B883C1894033FD97B883CBC62E3BEC13B8A3CBC62E3BEC13B8A3C7D6F7EBFA8FB8B3C7D6F7EBFA8FB8B3C574021BF8EBB8D3C574021BF8EBB8D3C865FA03E757B8F3C865FA03E757B8F3CBCE6773F5B3B913CBCE6773F5B3B913C5CB23B3F40FB923C5CB23B3F40FB923C2B4D34BE25BB943C2B4D34BE25BB943CC6676CBF0A7B963CC6676CBF0A7B963CA76252BFEE3A983CA76252BFEE3A983CD0FD103DD1FA993CD0FD103DD1FA993C822D5C3FB4BA9B3C822D5C3FB4BA9B3CF7DC643F977A9D3CF7DC643F977A9D3CFB0FD93D793A9F3CFB0FD93D793A9F3C178B47BF5BFAA03C178B47BF5BFAA03C9EC272BF3CBAA23C9EC272BF3CBAA23C56237BBE1D7AA43C56237BBE1D7AA43C3DEA2E3FFE39A63C3DEA2E3FFE39A63C68CC7B3FDDF9A73C68CC7B3FDDF9A73CF25BC23EBDB9A93CF25BC23EBDB9A93C28C912BF9C79AB3C28C912BF9C79AB3C05CC7FBF7A39AD3C05CC7FBF7A39AD3C2EA101BF58F9AE3C2EA101BF58F9AE3CED6FE73E35B9B03CED6FE73E35B9B03CF8AC7E3F1279B23CF8AC7E3F1279B23C2F7C1F3FEE38B43C2F7C1F3FEE38B43CB1ABA4BEC9F8B53CB1ABA4BEC9F8B53C007578BFA5B8B73C007578BFA5B8B73C03263ABF7F78B93C03263ABF7F78B93C6D373D3E5938BB3C6D373D3E5938BB3CFA436D3F33F8BC3CFA436D3F33F8BC3C0B16513F0BB8BE3C0B16513F0BB8BE3CCF3735BDE477C03CCF3735BDE477C03C3E535DBFBB37C23C3E535DBFBB37C23CC0D663BF93F7C33CC0D663BF93F7C33C4F07C7BD69B7C53C4F07C7BD69B7C53C7AF4483F3F77C73C7AF4483F3F77C73C0D08723F1537C93C0D08723F1537C93C5D57723EE9F6CA3C5D57723EE9F6CA3C0B9030BFBDB6CC3C0B9030BFBDB6CC3C38617BBF9176CE3C38617BBF9176CE3CAE28BEBE6436D03CAE28BEBE6436D03CEFA2143F36F6D13CEFA2143F36F6D13C5BB27F3F08B6D33C5BB27F3F08B6D33C5557FF3ED975D53C5557FF3ED975D53C7878EBBEA935D73C7878EBBEA935D73C58E57EBF79F5D83C58E57EBF79F5D83CD4B41DBF48B5DA3CD4B41DBF48B5DA3C8FF4A83E1775DC3C8FF4A83E1775DC3C49FE783FE534DE3C49FE783FE534DE3CED95383FB2F4DF3CED95383FB2F4DF3CE41D46BE7EB4E13CE41D46BE7EB4E13C6C1B6EBF4A74E33C6C1B6EBF4A74E33C3DC54FBF1534E53C3DC54FBF1534E53C2D6E593DE0F3E63C2D6E593DE0F3E63C8A745E3FA9B3E83C8A745E3FA9B3E83CF8CB623F7273EA3CF8CB623F7273EA3CA4FAB43D3B33EC3CA4FAB43D3B33EC3CD5594ABF02F3ED3CD5594ABF02F3ED3CA14871BFC9B2EF3CA14871BFC9B2EF3C888669BE9072F13C888669BE9072F13C4F32323F5532F33C4F32323F5532F33CFCF07A3F1AF2F43CFCF07A3F1AF2F43C99F1B93EDEB1F63C99F1B93EDEB1F63CBC7916BFA171F83CBC7916BFA171F83C90937FBF6431FA3C90937FBF6431FA3C2F67FBBE26F1FB3C2F67FBBE26F1FB3C4B7CEF3EE7B0FD3C4B7CEF3EE7B0FD3C9B187F3FA770FF3C9B187F3FA770FF3C4FEA1B3F3398003D4FEA1B3F3398003D093AADBE1378013D093AADBE1378013D938279BFF257023D938279BFF257023D240237BFD037033D240237BFD037033D61004F3EAE17043D61004F3EAE17043D18EE6E3F8CF7043D18EE6E3F8CF7043D45704E3F6AD7053D45704E3F6AD7053D2EA07DBD47B7063D2EA07DBD47B7063D60915FBF2397073D60915FBF2397073DA3BC61BFFF76083DA3BC61BFFF76083D59EAA2BDDB56093D59EAA2BDDB56093D21BB4B3FB7360A3D21BB4B3FB7360A3D5D84703F92160B3D5D84703F92160B3D04B1603E6CF60B3D04B1603E6CF60B3DFFD033BF46D60C3DFFD033BF46D60C3DB97B7ABF20B60D3DB97B7ABF20B60D3DC9B6B5BEF9950E3DC9B6B5BEF9950E3D844D183FD2750F3D844D183FD2750F3DA46F7F3FAB55103DA46F7F3FAB55103DFE71F73E8335113DFE71F73E8335113D507BF3BE5A15123D507BF3BE5A15123DC0467FBF31F5123DC0467FBF31F5123DAA1C1ABF08D5133DAA1C1ABF08D5133D0A7CB13EDEB4143D0A7CB13EDEB4143DDC017A3FB494153DDC017A3FB494153DAF6A353F8A74163DAF6A353F8A74163DB7DE57BE5E54173DB7DE57BE5E54173DF9BB6FBF3334183DF9BB6FBF3334183D28174DBF0714193D28174DBF0714193D8CE6903DDAF3193D8CE6903DDAF3193DB9A9603FAED31A3DB9A9603FAED31A3DC7A8603F80B31B3DC7A8603F80B31B3DC9D6903D52931C3DC9D6903D52931C3D57184DBF24731D3D57184DBF24731D3D47BB6FBFF5521E3D47BB6FBFF5521E3DFED657BEC6321F3DFED657BEC6321F3D146C353F9612203D146C353F9612203D6F017A3F66F2203D6F017A3F66F2203D5578B13E35D2213D5578B13E35D2213D3E1E1ABF04B2223D3E1E1ABF04B2223D9A467FBFD391233D9A467FBFD391233DD777F3BEA071243DD777F3BEA071243D7375F73E6E51253D7375F73E6E51253DC66F7F3F3B31263DC66F7F3F3B31263DED4B183F0711273DED4B183F0711273D7BBAB5BED3F0273D7BBAB5BED3F0273D217C7ABF9ED0283D217C7ABF9ED0283D97CF33BF69B0293D97CF33BF69B0293DB9B8603E34902A3DB9B8603E34902A3D0B85703FFD6F2B3D0B85703FFD6F2B3DEFB94B3FC74F2C3DEFB94B3FC74F2C3D1AFAA2BD902F2D3D1AFAA2BD902F2D3D92BD61BF580F2E3D92BD61BF580F2E3D69905FBF20EF2E3D69905FBF20EF2E3DA2807DBDE7CE2F3DA2807DBDE7CE2F3D70714E3FAEAE303D70714E3FAEAE303D62ED6E3F748E313D62ED6E3F748E313DA4F84E3E396E323DA4F84E3E396E323D860337BFFF4D333D860337BFFF4D333D228279BFC32D343D228279BFC32D343D5136ADBE870D353D5136ADBE870D353DE0EB1B3F4BED353DE0EB1B3F4BED353D70187F3F0ECD363D70187F3F0ECD363DCD78EF3ED0AC373DCD78EF3ED0AC373DA06AFBBE928C383DA06AFBBE928C383DAD937FBF536C393DAD937FBF536C393D237816BF144C3A3D237816BF144C3A3D47F5B93ED42B3B3D47F5B93ED42B3B3D61F17A3F940B3C3D61F17A3F940B3C3DE430323F53EB3C3DE430323F53EB3C3D3A8E69BE12CB3D3D3A8E69BE12CB3D3D4A4971BFD0AA3E3D4A4971BFD0AA3E3D9F584ABF8D8A3F3D9F584ABF8D8A3F3D620AB53D4A6A403D620AB53D4A6A403DE3CC623F064A413DE3CC623F064A413D90735E3FC229423D90735E3FC229423D9D4E593D7D09433D9D4E593D7D09433D65C64FBF37E9433D65C64FBF37E9433DB31A6EBFF1C8443DB31A6EBFF1C8443D231646BEABA8453D231646BEABA8453D4C97383F6388463D4C97383F6388463DD3FD783F1C68473DD3FD783F1C68473DD4F0A83ED347483DD4F0A83ED347483D62B61DBF8A27493D62B61DBF8A27493D29E57EBF40074A3D29E57EBF40074A3DF674EBBEF6E64A3DF674EBBEF6E64A3DC15AFF3EABC64B3DC15AFF3EABC64B3D73B27F3F60A64C3D73B27F3F60A64C3D53A1143F14864D3D53A1143F14864D3D592CBEBEC7654E3D592CBEBEC7654E3D97617BBF7A454F3D97617BBF7A454F3D9D8E30BF2C25503D9D8E30BF2C25503D0B5F723EDD04513D0B5F723EDD04513DB208723F8EE4513DB208723F8EE4513D40F3483F3EC4523D40F3483F3EC4523D0917C7BDEEA3533D0917C7BDEEA3533DA7D763BF9D83543DA7D763BF9D83543D40525DBF4B63553D40525DBF4B63553D3C1835BDF942563D3C1835BDF942563D2F17513FA622573D2F17513FA622573D3C436D3F5202583D3C436D3F5202583DA92F3D3EFEE1583DA92F3D3EFEE1583D5E273ABFA9C1593D5E273ABFA9C1593D867478BF53A15A3D867478BF53A15A3DF3A7A4BEFD805B3DF3A7A4BEFD805B3DBA7D1F3FA6605C3DBA7D1F3FA6605C3DC5AC7E3F4F405D3DC5AC7E3F4F405D3D666CE73EF61F5E3D666CE73EF61F5E3DE2A201BF9EFF5E3DE2A201BF9EFF5E3D"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x3727C5AC> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.reshape %arg20 : (tensor<256xi32>) -> tensor<256x1xi32> + %13 = "stablehlo.gather"(%arg0, %12) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<10x8xf32>, tensor<256x1xi32>) -> tensor<256x8xf32> + %14 = stablehlo.reshape %arg21 : (tensor<256xi32>) -> tensor<256x1xi32> + %15 = "stablehlo.gather"(%0, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %16 = "stablehlo.gather"(%1, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %17 = stablehlo.iota dim = 0 : tensor<256xi32> + %18 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<256xi32>) -> tensor<256x256xi32> + %19 = stablehlo.iota dim = 0 : tensor<256xi32> + %20 = stablehlo.broadcast_in_dim %19, dims = [1] : (tensor<256xi32>) -> tensor<256x256xi32> + %21 = stablehlo.compare LE, %20, %18, SIGNED : (tensor<256x256xi32>, tensor<256x256xi32>) -> tensor<256x256xi1> + %22 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x256xf32> + %23 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256x256xf32> + %24 = stablehlo.select %21, %22, %23 : tensor<256x256xi1>, tensor<256x256xf32> + %25 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %26 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %27 = stablehlo.multiply %13, %13 : tensor<256x8xf32> + %28 = stablehlo.reduce(%27 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %29 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %30 = stablehlo.divide %28, %29 : tensor<256xf32> + %31 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %32 = stablehlo.add %30, %31 : tensor<256xf32> + %33 = stablehlo.rsqrt %32 : tensor<256xf32> + %34 = stablehlo.broadcast_in_dim %33, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %35 = stablehlo.multiply %13, %34 : tensor<256x8xf32> + %36 = stablehlo.broadcast_in_dim %arg4, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %37 = stablehlo.multiply %35, %36 : tensor<256x8xf32> + %38 = stablehlo.dot_general %37, %arg9, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %39 = stablehlo.reshape %38 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %40 = stablehlo.dot_general %37, %arg7, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %41 = stablehlo.reshape %40 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %42 = stablehlo.dot_general %37, %arg10, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %43 = stablehlo.reshape %42 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %44 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %45 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %46 = stablehlo.multiply %39, %44 : tensor<256x2x4xf32> + %47 = stablehlo.slice %39 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %48 = stablehlo.slice %39 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %49 = stablehlo.negate %48 : tensor<256x2x2xf32> + %50 = stablehlo.concatenate %49, %47, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %51 = stablehlo.multiply %50, %45 : tensor<256x2x4xf32> + %52 = stablehlo.add %46, %51 : tensor<256x2x4xf32> + %53 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %54 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %55 = stablehlo.multiply %41, %53 : tensor<256x1x4xf32> + %56 = stablehlo.slice %41 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %57 = stablehlo.slice %41 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %58 = stablehlo.negate %57 : tensor<256x1x2xf32> + %59 = stablehlo.concatenate %58, %56, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %60 = stablehlo.multiply %59, %54 : tensor<256x1x4xf32> + %61 = stablehlo.add %55, %60 : tensor<256x1x4xf32> + %62 = stablehlo.reshape %61 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %63 = stablehlo.dynamic_update_slice %25, %62, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %64 = stablehlo.reshape %43 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %65 = stablehlo.dynamic_update_slice %26, %64, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %66 = stablehlo.reshape %52 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %67 = stablehlo.dot_general %66, %61, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %68 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %69 = stablehlo.multiply %67, %68 : tensor<1x256x2x256xf32> + %70 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %71 = stablehlo.add %69, %70 : tensor<1x256x2x256xf32> + %72 = stablehlo.reduce(%71 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %73 = stablehlo.broadcast_in_dim %72, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %74 = stablehlo.subtract %71, %73 : tensor<1x256x2x256xf32> + %75 = stablehlo.exponential %74 : tensor<1x256x2x256xf32> + %76 = stablehlo.reduce(%75 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %77 = stablehlo.broadcast_in_dim %76, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %78 = stablehlo.divide %75, %77 : tensor<1x256x2x256xf32> + %79 = stablehlo.dot_general %78, %43, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %80 = stablehlo.transpose %79, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %81 = stablehlo.reshape %80 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %82 = stablehlo.dot_general %81, %arg8, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %83 = stablehlo.add %13, %82 : tensor<256x8xf32> + %84 = stablehlo.multiply %83, %83 : tensor<256x8xf32> + %85 = stablehlo.reduce(%84 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %86 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %87 = stablehlo.divide %85, %86 : tensor<256xf32> + %88 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %89 = stablehlo.add %87, %88 : tensor<256xf32> + %90 = stablehlo.rsqrt %89 : tensor<256xf32> + %91 = stablehlo.broadcast_in_dim %90, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %92 = stablehlo.multiply %83, %91 : tensor<256x8xf32> + %93 = stablehlo.broadcast_in_dim %arg5, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %94 = stablehlo.multiply %92, %93 : tensor<256x8xf32> + %95 = stablehlo.dot_general %94, %arg3, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %96 = stablehlo.dot_general %94, %arg6, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %97 = stablehlo.negate %95 : tensor<256x16xf32> + %98 = stablehlo.exponential %97 : tensor<256x16xf32> + %99 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %100 = stablehlo.add %99, %98 : tensor<256x16xf32> + %101 = stablehlo.divide %99, %100 : tensor<256x16xf32> + %102 = stablehlo.multiply %95, %101 : tensor<256x16xf32> + %103 = stablehlo.multiply %102, %96 : tensor<256x16xf32> + %104 = stablehlo.dot_general %103, %arg2, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %105 = stablehlo.add %83, %104 : tensor<256x8xf32> + %106 = stablehlo.multiply %105, %105 : tensor<256x8xf32> + %107 = stablehlo.reduce(%106 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %108 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %109 = stablehlo.divide %107, %108 : tensor<256xf32> + %110 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %111 = stablehlo.add %109, %110 : tensor<256xf32> + %112 = stablehlo.rsqrt %111 : tensor<256xf32> + %113 = stablehlo.broadcast_in_dim %112, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %114 = stablehlo.multiply %105, %113 : tensor<256x8xf32> + %115 = stablehlo.broadcast_in_dim %arg13, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %116 = stablehlo.multiply %114, %115 : tensor<256x8xf32> + %117 = stablehlo.dot_general %116, %arg18, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %118 = stablehlo.reshape %117 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %119 = stablehlo.dot_general %116, %arg16, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %120 = stablehlo.reshape %119 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %121 = stablehlo.dot_general %116, %arg19, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %122 = stablehlo.reshape %121 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %123 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %124 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %125 = stablehlo.multiply %118, %123 : tensor<256x2x4xf32> + %126 = stablehlo.slice %118 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %127 = stablehlo.slice %118 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %128 = stablehlo.negate %127 : tensor<256x2x2xf32> + %129 = stablehlo.concatenate %128, %126, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %130 = stablehlo.multiply %129, %124 : tensor<256x2x4xf32> + %131 = stablehlo.add %125, %130 : tensor<256x2x4xf32> + %132 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %133 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %134 = stablehlo.multiply %120, %132 : tensor<256x1x4xf32> + %135 = stablehlo.slice %120 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %136 = stablehlo.slice %120 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %137 = stablehlo.negate %136 : tensor<256x1x2xf32> + %138 = stablehlo.concatenate %137, %135, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %139 = stablehlo.multiply %138, %133 : tensor<256x1x4xf32> + %140 = stablehlo.add %134, %139 : tensor<256x1x4xf32> + %141 = stablehlo.reshape %140 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %142 = stablehlo.dynamic_update_slice %63, %141, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %143 = stablehlo.reshape %122 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %144 = stablehlo.dynamic_update_slice %65, %143, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %145 = stablehlo.reshape %131 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %146 = stablehlo.dot_general %145, %140, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %147 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %148 = stablehlo.multiply %146, %147 : tensor<1x256x2x256xf32> + %149 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %150 = stablehlo.add %148, %149 : tensor<1x256x2x256xf32> + %151 = stablehlo.reduce(%150 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %152 = stablehlo.broadcast_in_dim %151, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %153 = stablehlo.subtract %150, %152 : tensor<1x256x2x256xf32> + %154 = stablehlo.exponential %153 : tensor<1x256x2x256xf32> + %155 = stablehlo.reduce(%154 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %156 = stablehlo.broadcast_in_dim %155, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %157 = stablehlo.divide %154, %156 : tensor<1x256x2x256xf32> + %158 = stablehlo.dot_general %157, %122, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %159 = stablehlo.transpose %158, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %160 = stablehlo.reshape %159 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %161 = stablehlo.dot_general %160, %arg17, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %162 = stablehlo.add %105, %161 : tensor<256x8xf32> + %163 = stablehlo.multiply %162, %162 : tensor<256x8xf32> + %164 = stablehlo.reduce(%163 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %165 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %166 = stablehlo.divide %164, %165 : tensor<256xf32> + %167 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %168 = stablehlo.add %166, %167 : tensor<256xf32> + %169 = stablehlo.rsqrt %168 : tensor<256xf32> + %170 = stablehlo.broadcast_in_dim %169, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %171 = stablehlo.multiply %162, %170 : tensor<256x8xf32> + %172 = stablehlo.broadcast_in_dim %arg14, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %173 = stablehlo.multiply %171, %172 : tensor<256x8xf32> + %174 = stablehlo.dot_general %173, %arg12, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %175 = stablehlo.dot_general %173, %arg15, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %176 = stablehlo.negate %174 : tensor<256x16xf32> + %177 = stablehlo.exponential %176 : tensor<256x16xf32> + %178 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %179 = stablehlo.add %178, %177 : tensor<256x16xf32> + %180 = stablehlo.divide %178, %179 : tensor<256x16xf32> + %181 = stablehlo.multiply %174, %180 : tensor<256x16xf32> + %182 = stablehlo.multiply %181, %175 : tensor<256x16xf32> + %183 = stablehlo.dot_general %182, %arg11, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %184 = stablehlo.add %162, %183 : tensor<256x8xf32> + %185 = stablehlo.multiply %184, %184 : tensor<256x8xf32> + %186 = stablehlo.reduce(%185 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %187 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %188 = stablehlo.divide %186, %187 : tensor<256xf32> + %189 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %190 = stablehlo.add %188, %189 : tensor<256xf32> + %191 = stablehlo.rsqrt %190 : tensor<256xf32> + %192 = stablehlo.broadcast_in_dim %191, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %193 = stablehlo.multiply %184, %192 : tensor<256x8xf32> + %194 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %195 = stablehlo.multiply %193, %194 : tensor<256x8xf32> + %196 = stablehlo.constant dense<1> : tensor + %197 = stablehlo.subtract %arg22, %196 : tensor + %198 = stablehlo.dynamic_slice %195, %197, %9, sizes = [1, 8] : (tensor<256x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %199 = stablehlo.reshape %198 : (tensor<1x8xf32>) -> tensor<8xf32> + %200 = stablehlo.dot_general %199, %arg0, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %201 = stablehlo.constant dense<0> : tensor + %202 = stablehlo.constant dense<0xFF800000> : tensor + %203 = stablehlo.iota dim = 0 : tensor<10xi32> + %204:2 = stablehlo.reduce(%200 init: %202), (%203 init: %201) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %205 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %206 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %207 = stablehlo.or %205, %206 : tensor + %208 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %209 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %210 = stablehlo.and %208, %209 : tensor + %211 = stablehlo.or %207, %210 : tensor + %212 = stablehlo.select %207, %amv_l, %amv_r : tensor, tensor + %213 = stablehlo.select %211, %ami_l, %ami_r : tensor, tensor + stablehlo.return %212, %213 : tensor, tensor + } + return %204#1, %142, %144 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/internlm3/README.md b/src/lib/mlxcel-xla/assets/internlm3/README.md new file mode 100644 index 000000000..c8a259987 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/internlm3/README.md @@ -0,0 +1,19 @@ +# InternLM3 config fixture (#499 dense arch pack) + +`config.json` is a **small synthetic** InternLM3 config (real switches, tiny dims) +registered as a byte-exact structural fixture in `validation.rs`; the real +InternLM3-8B `config.json` parsing is covered by the `config::tests` parse assertion. + +InternLM3 (`model_type` `internlm3`) is the proven **Llama** forward with standard +names: **untied** embeddings, no bias (`qkv_bias = false`; a `qkv_bias = true` +checkpoint turns the q/k/v bias on), and `rope_type = "dynamic"` served as plain +RoPE. Dynamic NTK RoPE is identity within the original context and only rescales +beyond it, so short / in-context generation is served with plain RoPE (`rope_theta`); +the long-context NTK rescale is a follow-up. Verified from the checkpoint's +`modeling_internlm3.py` to use the emitter's half-split RoPE, `cat(freqs, freqs)` +table and `head_dim^-0.5` scaling. + +Validation: `validation::dense_pack_families_reuse_proven_graphs` asserts it emits +byte-for-byte identical StableHLO to an untied plain-RoPE Llama reference (the +execution-proven forward). The full token-exact / serve gate runs post-merge on the +real checkpoint via `scripts/xla/validate_arch.sh`. diff --git a/src/lib/mlxcel-xla/assets/internlm3/config.json b/src/lib/mlxcel-xla/assets/internlm3/config.json new file mode 100644 index 000000000..e815609e5 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/internlm3/config.json @@ -0,0 +1,16 @@ +{ + "model_type": "internlm3", + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 16, + "num_hidden_layers": 2, + "rms_norm_eps": 1e-05, + "rope_theta": 50000000.0, + "vocab_size": 10, + "tie_word_embeddings": false, + "qkv_bias": false, + "rope_scaling": { "rope_type": "dynamic", "factor": 6.0 }, + "hidden_act": "silu" +} diff --git a/src/lib/mlxcel-xla/assets/internlm3/decode.mlir b/src/lib/mlxcel-xla/assets/internlm3/decode.mlir new file mode 100644 index 000000000..3e83c4e13 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/internlm3/decode.mlir @@ -0,0 +1,212 @@ +module @decode_step { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg4: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg5: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg6: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg7: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg10: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg11: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg12: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg13: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg14: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg15: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg16: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg17: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg18: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg19: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg20: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg21: tensor loc("token"), %arg22: tensor loc("pos"), %arg23: tensor loc("cache_len"), %arg24: tensor<2x256x1x4xf32> loc("kcache"), %arg25: tensor<2x256x1x4xf32> loc("vcache")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3F0000803F40510A3F0000803F3311D5BEFFFF7F3F3311D5BEFFFF7F3F26707DBFFEFF7F3F26707DBFFEFF7F3F305527BFFDFF7F3F305527BFFDFF7F3F2C3C913EFCFF7F3F2C3C913EFCFF7F3FB8CD753FFAFF7F3FB8CD753FFAFF7F3FBDFF403FF8FF7F3FBDFF403FF8FF7F3FF6FD14BEF5FF7F3FF6FD14BEF5FF7F3FD53F69BFF2FF7F3FD53F69BFF2FF7F3F64CD56BFEFFF7F3F64CD56BFEFFF7F3F7205913BECFF7F3F7205913BECFF7F3FD006583FE8FF7F3FD006583FE8FF7F3F6F4E683FE4FF7F3F6F4E683FE4FF7F3FD7040C3EDFFF7F3FD7040C3EDFFF7F3FE87A42BFDAFF7F3FE87A42BFDAFF7F3F2C2975BFD5FF7F3F2C2975BFD5FF7F3F36E28CBED0FF7F3F36E28CBED0FF7F3F840A293FCAFF7F3F840A293FCAFF7F3FBF1B7D3FC3FF7F3FBF1B7D3FC3FF7F3F22F0D03EBDFF7F3F22F0D03EBDFF7F3FFC370CBFB6FF7F3FFC370CBFB6FF7F3F6FFD7FBFAFFF7F3F6FFD7FBFAFFF7F3FBF6708BFA7FF7F3FBF6708BFA7FF7F3FFE2DD93E9FFF7F3FFE2DD93E9FFF7F3F78BF7D3F97FF7F3F78BF7D3F97FF7F3F819C253F8FFF7F3F819C253F8FFF7F3F389395BE86FF7F3F389395BE86FF7F3F576D76BF7CFF7F3F576D76BF7CFF7F3FB3803FBF73FF7F3FB3803FBF73FF7F3F18F41D3E69FF7F3F18F41D3E69FF7F3F8E2C6A3F5FFF7F3F8E2C6A3F5FFF7F3FAA8F553F54FF7F3FAA8F553F54FF7F3FB78659BC49FF7F3FB78659BC49FF7F3FE73B59BF3EFF7F3FE73B59BF3EFF7F3F5F5867BF32FF7F3F5F5867BF32FF7F3FEA0803BE27FF7F3FEA0803BE27FF7F3F2DF2433F1AFF7F3F2DF2433F1AFF7F3FB57F743F0EFF7F3FB57F743F0EFF7F3F6C85883E01FF7F3F6C85883E01FF7F3F74BC2ABFF4FE7F3F74BC2ABFF4FE7F3F44C27CBFE6FE7F3F44C27CBFE6FE7F3FE0CACCBED8FE7F3FE0CACCBED8FE7F3FE81B0E3FCAFE7F3FE81B0E3FCAFE7F3FBBF57F3FBBFE7F3FBBF57F3FBBFE7F3F807B063FACFE7F3F807B063FACFE7F3F6D46DDBE9DFE7F3F6D46DDBE9DFE7F3FB3097EBF8DFE7F3FB3097EBF8DFE7F3F80E023BF7DFE7F3F80E023BF7DFE7F3F44E7993E6DFE7F3F44E7993E6DFE7F3F0308773F5DFE7F3F0308773F5DFE7F3FD1FD3D3F4CFE7F3FD1FD3D3F4CFE7F3F0EE726BE3AFE7F3F0EE726BE3AFE7F3F95146BBF29FE7F3F95146BBF29FE7F3FA64D54BF17FE7F3FA64D54BF17FE7F3F2C43B53C04FE7F3F2C43B53C04FE7F3FA26C5A3FF2FD7F3FA26C5A3FF2FD7F3FAC5D663FDFFD7F3FAC5D663FDFFD7F3FB714F43DCCFD7F3FB714F43DCCFD7F3F836545BFB8FD7F3F836545BFB8FD7F3F56D173BFA4FD7F3F56D173BFA4FD7F3FE62584BE90FD7F3FE62584BE90FD7F3FF76A2C3F7BFD7F3FF76A2C3F7BFD7F3FB8637C3F66FD7F3FB8637C3F66FD7F3F83A1C83E51FD7F3F83A1C83E51FD7F3FFAFC0FBF3BFD7F3FFAFC0FBF3BFD7F3FE5E87FBF25FD7F3FE5E87FBF25FD7F3F908C04BF0FFD7F3F908C04BF0FFD7F3F6D5AE13EF8FC7F3F6D5AE13EF8FC7F3FD54E7E3FE1FC7F3FD54E7E3FE1FC7F3F3521223FCAFC7F3F3521223FCAFC7F3F3A389EBEB2FC7F3F3A389EBEB2FC7F3FBC9D77BF9AFC7F3FBC9D77BF9AFC7F3F20773CBF82FC7F3F20773CBF82FC7F3FACD62F3E69FC7F3FACD62F3E69FC7F3FE4F76B3F50FC7F3FE4F76B3F50FC7F3F6107533F37FC7F3F6107533F37FC7F3F5ABFFDBC1DFC7F3F5ABFFDBC1DFC7F3FFC985BBF03FC7F3FFC985BBF03FC7F3F5A5E65BFE9FB7F3F5A5E65BFE9FB7F3FB512E2BDCEFB7F3FB512E2BDCEFB7F3FE4D4463FB3FB7F3FE4D4463FB3FB7F3F141E733F98FB7F3F141E733F98FB7F3F72877F3E7CFB7F3F72877F3E7CFB7F3F05162EBF60FB7F3F05162EBF60FB7F3F1C007CBF44FB7F3F1C007CBF44FB7F3F1F74C4BE27FB7F3F1F74C4BE27FB7F3F28DB113F0AFB7F3F28DB113F0AFB7F3FEDD67F3FEDFA7F3FEDD67F3FEDFA7F3FF69A023FCFFA7F3FF69A023FCFFA7F3FE869E5BEB1FA7F3FE869E5BEB1FA7F3FDE8E7EBF93FA7F3FDE8E7EBF93FA7F3FAA5E20BF74FA7F3FAA5E20BF74FA7F3F0486A23E55FA7F3F0486A23E55FA7F3F7D2E783F36FA7F3F7D2E783F36FA7F3FA8EC3A3F16FA7F3FA8EC3A3F16FA7F3FC2C238BEF6F97F3FC2C238BEF6F97F3F78D66CBFD5F97F3F78D66CBFD5F97F3FE0BC51BFB5F97F3FE0BC51BFB5F97F3F381B233D94F97F3F381B233D94F97F3FEEC05C3F72F97F3FEEC05C3F72F97F3F6E5A643F51F97F3F6E5A643F51F97F3F2A0CD03D2FF97F3F2A0CD03D2FF97F3F494048BF0CF97F3F494048BF0CF97F3FF16572BFE9F87F3FF16572BFE9F87F3FF8BD76BEC6F87F3FF8BD76BEC6F87F3F95BD2F3FA3F87F3F95BD2F3FA3F87F3F71977B3F7FF87F3F71977B3F7FF87F3FCB42C03E5BF87F3FCB42C03E5BF87F3F6AB613BF37F87F3F6AB613BF37F87F3FD4BF7FBF12F87F3FD4BF7FBF12F87F3FBFA600BFEDF77F3FBFA600BFEDF77F3FC874E93EC8F77F3FC874E93EC8F77F3FCBC97E3FA2F77F3FCBC97E3FA2F77F3FE7981E3F7CF77F3FE7981E3F7CF77F3F8BD0A6BE55F77F3F8BD0A6BE55F77F3F44BA78BF2FF77F3F44BA78BF2FF77F3F6F5E39BF07F77F3F6F5E39BF07F77F3F25AB413EE0F67F3F25AB413EE0F67F3F4CB06D3FB8F67F3F4CB06D3FB8F67F3F2A6E503F90F67F3F2A6E503F90F67F3F7E5347BD68F67F3F7E5347BD68F67F3F73E45DBF3FF67F3F73E45DBF3FF67F3FEE5163BF16F67F3FEE5163BF16F67F3F7301BEBDECF57F3F7301BEBDECF57F3FA9A7493FC3F57F3FA9A7493FC3F57F3FF2A8713F99F57F3FF2A8713F99F57F3F8BEF6D3E6EF57F3F8BEF6D3E6EF57F3F9F6131BF43F57F3F9F6131BF43F57F3FBB297BBF18F57F3FBB297BBF18F57F3F9C0DBCBEEDF47F3F9C0DBCBEEDF47F3FB68E153FC1F47F3FB68E153FC1F47F3F99A37F3F95F47F3F99A37F3F95F47F3FE45FFD3E68F47F3FE45FFD3E68F47F3FFA7AEDBE3CF47F3FFA7AEDBE3CF47F3F9DFF7EBF0EF47F3F9DFF7EBF0EF47F3FF6CF1CBFE1F37F3FF6CF1CBFE1F37F3FBA17AB3EB3F37F3FBA17AB3EB3F37F3F0E41793F85F37F3F0E41793F85F37F3F7FCC373F57F37F3F7FCC373F57F37F3FA48F4ABE28F37F3FA48F4ABE28F37F3F5B856EBFF9F27F3F5B856EBFF9F27F3F461B4FBFC9F27F3F461B4FBFC9F27F3FC5876B3D99F27F3FC5876B3D99F27F3F84035F3F69F27F3F84035F3F69F27F3FDF44623F39F27F3FDF44623F39F27F3FEDF2AB3D08F27F3FEDF2AB3D08F27F3FFD0A4BBFD7F17F3FFD0A4BBFD7F17F3F1AE770BFA5F17F3F1AE770BFA5F17F3F591C65BE73F17F3F591C65BE73F17F3F1A02333F41F17F3F1A02333F41F17F3FFCB67A3F0FF17F3FFCB67A3F0FF17F3FA7D4B73EDCF07F3FA7D4B73EDCF07F3F016417BFA9F07F3F016417BFA9F07F3F3E827FBF75F07F3F3E827FBF75F07F3F366DF9BE41F07F3F366DF9BE41F07F3F697CF13E0DF07F3F697CF13E0DF07F3F51307F3FD9EF7F3F51307F3FD9EF7F3FE0031B3FA4EF7F3FE0031B3FA4EF7F3F7A5BAFBE6FEF7F3F7A5BAFBE6FEF7F3FD8C279BF39EF7F3FD8C279BF39EF7F3FDF3636BF03EF7F3FDF3636BF03EF7F3F1470533ECDEE7F3F1470533ECDEE7F3FA2556F3F97EE7F3FA2556F3F97EE7F3F3BC44D3F60EE7F3F3BC44D3F60EE7F3FA9DB87BD29EE7F3FA9DB87BD29EE7F3F1C1E60BFF1ED7F3F1C1E60BFF1ED7F3F463361BFB9ED7F3F463361BFB9ED7F3FF3E099BD81ED7F3FF3E099BD81ED7F3F3F6A4C3F48ED7F3F3F6A4C3F48ED7F3F6D20703F10ED7F3F6D20703F10ED7F3F8E445C3ED6EC7F3F8E445C3ED6EC7F3FFE9E34BF9DEC7F3FFE9E34BF9DEC7F3F353F7ABF63EC7F3F353F7ABF63EC7F3F0298B3BE29EC7F3F0298B3BE29EC7F3F4336193FEEEB7F3F4336193FEEEB7F3FC35B7F3FB3EB7F3FC35B7F3FB3EB7F3F8875F53E78EB7F3F8875F53E78EB7F3FFF78F5BE3DEB7F3FFF78F5BE3DEB7F3FE75B7FBF01EB7F3FE75B7FBF01EB7F3FAE3419BFC4EA7F3FAE3419BFC4EA7F3FB69BB33E88EA7F3FB69BB33E88EA7F3F9F3F7A3F4BEA7F3F9F3F7A3F4BEA7F3F979D343F0EEA7F3F979D343F0EEA7F3F464C5CBED0E97F3F464C5CBED0E97F3F1C2170BF92E97F3F1C2170BF92E97F3F0F694CBF54E97F3F0F694CBF54E97F3FB6F0993D16E97F3FB6F0993D16E97F3F3634613FD7E87F3F3634613FD7E87F3F281D603F97E87F3F281D603F97E87F3FE4CB873D58E87F3FE4CB873D58E87F3F67C54DBF18E87F3F67C54DBF18E87F3FEF546FBFD8E77F3FEF546FBFD8E77F3F586853BE97E77F3F586853BE97E77F3F4238363F56E77F3F4238363F56E77F3F69C2793F15E77F3F69C2793F15E77F3FC357AF3ED3E67F3FC357AF3ED3E67F3F72051BBF91E67F3F72051BBF91E67F3F29307FBF4FE67F3F29307FBF4FE67F3FED78F1BE0CE67F3FED78F1BE0CE67F3FAA70F93ECAE57F3FAA70F93ECAE57F3F5D827F3F86E57F3F5D827F3F86E57F3F6962173F43E57F3F6962173F43E57F3F57D8B7BEFFE47F3F57D8B7BEFFE47F3F62B77ABFBAE47F3F62B77ABFBAE47F3FB00033BF76E47F3FB00033BF76E47F3F0D24653E31E47F3F0D24653E31E47F3FC5E7703FECE37F3FC5E7703FECE37F3FC9094B3FA6E37F3FC9094B3FA6E37F3FAC02ACBD60E37F3FAC02ACBD60E37F3FCB4562BF1AE37F3FCB4562BF1AE37F3F8C025FBFD3E27F3F8C025FBFD3E27F3F36686BBD8CE27F3F36686BBD8CE27F3F6F1C4F3F45E27F3F6F1C4F3F45E27F3FA4846E3FFDE17F3FA4846E3FFDE17F3FE5874A3EB5E17F3FE5874A3EB5E17F3FDFCD37BF6DE17F3FDFCD37BF6DE17F3F9A4079BF24E17F3F9A4079BF24E17F3F0014ABBEDBE07F3F0014ABBEDBE07F3F86D11C3F92E07F3F86D11C3F92E07F3F70FF7E3F48E07F3F70FF7E3F48E07F3F7A77ED3EFFDF7F3F7A77ED3EFFDF7F3F5363FDBEB4DF7F3F5363FDBEB4DF7F3FB4A37FBF6ADF7F3FB4A37FBF6ADF7F3F1B8D15BF1FDF7F3F1B8D15BF1FDF7F3F4911BC3ED3DE7F3F4911BC3ED3DE7F3F1D2A7B3F88DE7F3F1D2A7B3F88DE7F3F3260313F3CDE7F3F3260313F3CDE7F3F3BF76DBEEFDD7F3F3BF76DBEEFDD7F3F99A971BFA3DD7F3F99A971BFA3DD7F3F71A649BF56DD7F3F71A649BF56DD7F3F3011BE3D08DD7F3F3011BE3D08DD7F3FD752633FBBDC7F3FD752633FBBDC7F3F76E35D3F6DDC7F3F76E35D3F6DDC7F3FEC33473D1EDC7F3FEC33473D1EDC7F3F506F50BFD0DB7F3F506F50BFD0DB7F3F90AF6DBF81DB7F3F90AF6DBF81DB7F3F62A341BE31DB7F3F62A341BE31DB7F3FCC5F393FE2DA7F3FCC5F393FE2DA7F3FCCB9783F92DA7F3FCCB9783F92DA7F3FCFCCA63E41DA7F3FCFCCA63E41DA7F3F749A1EBFF1D97F3F749A1EBFF1D97F3F9AC97EBFA0D97F3F9AC97EBFA0D97F3F4471E9BE4ED97F3F4471E9BE4ED97F3F74A8003FFDD87F3F74A8003FFDD87F3FEABF7F3FAAD87F3FEABF7F3FAAD87F3FCDB4133F58D87F3FCDB4133F58D87F3F7546C0BE05D87F3F7546C0BE05D87F3FCF977BBFB2D77F3FCF977BBFB2D77F3F25BC2FBF5FD77F3F25BC2FBF5FD77F3FA3C5763E0BD77F3FA3C5763E0BD77F3F9466723FB7D67F3F9466723FB7D67F3F0E3F483F63D67F3F0E3F483F63D67F3FE31BD0BD0ED67F3FE31BD0BD0ED67F3F535B64BFB9D57F3F535B64BFB9D57F3FEEBF5CBF64D57F3FEEBF5CBF64D57F3F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F824A1439A46A573F824A1439B7C7683F814A9439B7C7683F814A9439C381103EC26FDE39C381103EC26FDE39CFBD41BF814A143ACFBD41BF814A143A107C75BF215D393A107C75BF215D393A8C0F8FBEC16F5E3A8C0F8FBEC16F5E3A4630283F30C1813A4630283F30C1813A95467D3F804A943A95467D3F804A943A3201D33ECFD3A63A3201D33ECFD3A63AF8440BBF1E5DB93AF8440BBF1E5DB93A5CFF7FBF6DE6CB3A5CFF7FBF6DE6CB3AD85C09BFBB6FDE3AD85C09BFBB6FDE3A2220D73E0AF9F03A2220D73E0AF9F03A72987D3F2CC1013B72987D3F2CC1013B4479263FD3050B3B4479263FD3050B3B106893BE794A143B106893BE794A143B251E76BF208F1D3B251E76BF208F1D3BB34040BFC6D3263BB34040BFC6D3263B6979193E6C18303B6979193E6C18303BC8B6693F125D393BC8B6693F125D393B102F563FB7A1423B102F563FB7A1423B150511BC5DE64B3B150511BC5DE64B3BE7A158BF022B553BE7A158BF022B553BFCD367BFA66F5E3BFCD367BFA66F5E3B388707BE4BB4673B388707BE4BB4673B0837433FEFF8703B0837433FEFF8703B0ED5743F933D7A3B0ED5743F933D7A3B2AB48A3E1BC1813B2AB48A3E1BC1813BE9E329BF6D63863BE9E329BF6D63863BA4EF7CBFBE058B3BA4EF7CBFBE058B3B06DECEBE0FA88F3B06DECEBE0FA88F3B4D2A0D3F604A943B4D2A0D3F604A943B39FA7F3FB1EC983B39FA7F3FB1EC983BF771073F028F9D3BF771073F028F9D3BC23ADBBE5231A23BC23ADBBE5231A23B38E57DBFA3D3A63B38E57DBFA3D3A63BEABE24BFF375AB3BEABE24BFF375AB3B9FBD973E4218B03B9FBD973E4218B03B4BBB763F92BAB43B4BBB763F92BAB43BBCBF3E3FE15CB93BBCBF3E3FE15CB93BFB6D22BE30FFBD3BFB6D22BE30FFBD3B28A16ABF7FA1C23B28A16ABF7FA1C23B31EF54BFCE43C73B31EF54BFCE43C73BA103913C1CE6CB3BA103913C1CE6CB3BD0D4593F6A88D03BD0D4593F6A88D03B9ADB663FB82AD53B9ADB663FB82AD53BE813FD3D05CDD93BE813FD3D05CDD93B56AC44BF526FDE3B56AC44BF526FDE3B222974BF9F11E33B222974BF9F11E33BFF5586BEECB3E73BFF5586BEECB3E73B24942B3F3856EC3B24942B3F3856EC3BA0937C3F84F8F03BA0937C3F84F8F03BB4B6CA3ED09AF53BB4B6CA3ED09AF53BCD0C0FBF1B3DFA3BCD0C0FBF1B3DFA3BF4EF7FBF66DFFE3BF4EF7FBF66DFFE3B5E8405BFD9C0013C5E8405BFD9C0013CFD50DF3EFE11043CFD50DF3EFE11043CE72C7E3F2363063CE72C7E3F2363063C4301233F48B4083C4301233F48B4083C23109CBE6C050B3C23109CBE6C050B3C7F5377BF91560D3C7F5377BF91560D3CF23A3DBFB5A70F3CF23A3DBFB5A70F3C4B5F2B3ED9F8113C4B5F2B3ED9F8113CD4866B3FFD49143CD4866B3FFD49143C0BAB533F219B163C0BAB533F219B163CCE81D9BC44EC183CCE81D9BC44EC183C5B035BBF673D1B3C5B035BBF673D1B3C97DE65BF8B8E1D3C97DE65BF8B8E1D3C4D14EBBDADDF1F3C4D14EBBDADDF1F3CB31D463FD030223CB31D463FD030223C5178733FF381243C5178733FF381243C23F5813E15D3263C23F5813E15D3263CED402DBF3724293CED402DBF3724293C8C327CBF59752B3C8C327CBF59752B3C518BC6BE7AC62D3C518BC6BE7AC62D3C6EEC103F9C17303C6EEC103F9C17303C8EE07F3FBD68323C8EE07F3FBD68323C1894033FDEB9343C1894033FDEB9343CBC62E3BEFE0A373CBC62E3BEFE0A373C7D6F7EBF1F5C393C7D6F7EBF1F5C393C574021BF3FAD3B3C574021BF3FAD3B3C865FA03E5FFE3D3C865FA03E5FFE3D3CBCE6773F7F4F403CBCE6773F7F4F403C5CB23B3F9EA0423C5CB23B3F9EA0423C2B4D34BEBDF1443C2B4D34BEBDF1443CC6676CBFDC42473CC6676CBFDC42473CA76252BFFB93493CA76252BFFB93493CD0FD103D19E54B3CD0FD103D19E54B3C822D5C3F37364E3C822D5C3F37364E3CF7DC643F5587503CF7DC643F5587503CFB0FD93D73D8523CFB0FD93D73D8523C178B47BF9029553C178B47BF9029553C9EC272BFAD7A573C9EC272BFAD7A573C56237BBECACB593C56237BBECACB593C3DEA2E3FE61C5C3C3DEA2E3FE61C5C3C68CC7B3F036E5E3C68CC7B3F036E5E3CF25BC23E1EBF603CF25BC23E1EBF603C28C912BF3A10633C28C912BF3A10633C05CC7FBF5561653C05CC7FBF5561653C2EA101BF70B2673C2EA101BF70B2673CED6FE73E8B036A3CED6FE73E8B036A3CF8AC7E3FA5546C3CF8AC7E3FA5546C3C2F7C1F3FBFA56E3C2F7C1F3FBFA56E3CB1ABA4BED9F6703CB1ABA4BED9F6703C007578BFF347733C007578BFF347733C03263ABF0C99753C03263ABF0C99753C6D373D3E25EA773C6D373D3E25EA773CFA436D3F3D3B7A3CFA436D3F3D3B7A3C0B16513F558C7C3C0B16513F558C7C3CCF3735BD6DDD7E3CCF3735BD6DDD7E3C3E535DBF4297803C3E535DBF4297803CC0D663BFCEBF813CC0D663BFCEBF813C4F07C7BD59E8823C4F07C7BD59E8823C7AF4483FE410843C7AF4483FE410843C0D08723F7039853C0D08723F7039853C5D57723EFA61863C5D57723EFA61863C0B9030BF858A873C0B9030BF858A873C38617BBF10B3883C38617BBF10B3883CAE28BEBE9ADB893CAE28BEBE9ADB893CEFA2143F24048B3CEFA2143F24048B3C5BB27F3FAE2C8C3C5BB27F3FAE2C8C3C5557FF3E38558D3C5557FF3E38558D3C7878EBBEC27D8E3C7878EBBEC27D8E3C58E57EBF4BA68F3C58E57EBF4BA68F3CD4B41DBFD4CE903CD4B41DBFD4CE903C8FF4A83E5DF7913C8FF4A83E5DF7913C49FE783FE61F933C49FE783FE61F933CED95383F6F48943CED95383F6F48943CE41D46BEF770953CE41D46BEF770953C6C1B6EBF8099963C6C1B6EBF8099963C3DC54FBF08C2973C3DC54FBF08C2973C2D6E593D90EA983C2D6E593D90EA983C8A745E3F17139A3C8A745E3F17139A3CF8CB623F9F3B9B3CF8CB623F9F3B9B3CA4FAB43D26649C3CA4FAB43D26649C3CD5594ABFAD8C9D3CD5594ABFAD8C9D3CA14871BF34B59E3CA14871BF34B59E3C888669BEBBDD9F3C888669BEBBDD9F3C4F32323F4106A13C4F32323F4106A13CFCF07A3FC72EA23CFCF07A3FC72EA23C99F1B93E4D57A33C99F1B93E4D57A33CBC7916BFD37FA43CBC7916BFD37FA43C90937FBF59A8A53C90937FBF59A8A53C2F67FBBEDED0A63C2F67FBBEDED0A63C4B7CEF3E63F9A73C4B7CEF3E63F9A73C9B187F3FE821A93C9B187F3FE821A93C4FEA1B3F6D4AAA3C4FEA1B3F6D4AAA3C093AADBEF172AB3C093AADBEF172AB3C938279BF769BAC3C938279BF769BAC3C240237BFFAC3AD3C240237BFFAC3AD3C61004F3E7EECAE3C61004F3E7EECAE3C18EE6E3F0115B03C18EE6E3F0115B03C45704E3F843DB13C45704E3F843DB13C2EA07DBD0866B23C2EA07DBD0866B23C60915FBF8B8EB33C60915FBF8B8EB33CA3BC61BF0DB7B43CA3BC61BF0DB7B43C59EAA2BD90DFB53C59EAA2BD90DFB53C21BB4B3F1208B73C21BB4B3F1208B73C5D84703F9430B83C5D84703F9430B83C04B1603E1559B93C04B1603E1559B93CFFD033BF9781BA3CFFD033BF9781BA3CB97B7ABF18AABB3CB97B7ABF18AABB3CC9B6B5BE99D2BC3CC9B6B5BE99D2BC3C844D183F1AFBBD3C844D183F1AFBBD3CA46F7F3F9A23BF3CA46F7F3F9A23BF3CFE71F73E1A4CC03CFE71F73E1A4CC03C507BF3BE9A74C13C507BF3BE9A74C13CC0467FBF1A9DC23CC0467FBF1A9DC23CAA1C1ABF9AC5C33CAA1C1ABF9AC5C33C0A7CB13E19EEC43C0A7CB13E19EEC43CDC017A3F9816C63CDC017A3F9816C63CAF6A353F163FC73CAF6A353F163FC73CB7DE57BE9567C83CB7DE57BE9567C83CF9BB6FBF1390C93CF9BB6FBF1390C93C28174DBF91B8CA3C28174DBF91B8CA3C8CE6903D0EE1CB3C8CE6903D0EE1CB3CB9A9603F8C09CD3CB9A9603F8C09CD3CC7A8603F0932CE3CC7A8603F0932CE3CC9D6903D865ACF3CC9D6903D865ACF3C57184DBF0283D03C57184DBF0283D03C47BB6FBF7FABD13C47BB6FBF7FABD13CFED657BEFBD3D23CFED657BEFBD3D23C146C353F76FCD33C146C353F76FCD33C6F017A3FF224D53C6F017A3FF224D53C5578B13E6D4DD63C5578B13E6D4DD63C3E1E1ABFE875D73C3E1E1ABFE875D73C9A467FBF629ED83C9A467FBF629ED83CD777F3BEDDC6D93CD777F3BEDDC6D93C7375F73E57EFDA3C7375F73E57EFDA3CC66F7F3FD117DC3CC66F7F3FD117DC3CED4B183F4A40DD3CED4B183F4A40DD3C7BBAB5BEC368DE3C7BBAB5BEC368DE3C217C7ABF3C91DF3C217C7ABF3C91DF3C97CF33BFB5B9E03C97CF33BFB5B9E03CB9B8603E2DE2E13CB9B8603E2DE2E13C0B85703FA50AE33C0B85703FA50AE33CEFB94B3F1D33E43CEFB94B3F1D33E43C1AFAA2BD945BE53C1AFAA2BD945BE53C92BD61BF0B84E63C92BD61BF0B84E63C69905FBF82ACE73C69905FBF82ACE73CA2807DBDF8D4E83CA2807DBDF8D4E83C70714E3F6FFDE93C70714E3F6FFDE93C62ED6E3FE425EB3C62ED6E3FE425EB3CA4F84E3E5A4EEC3CA4F84E3E5A4EEC3C860337BFCF76ED3C860337BFCF76ED3C228279BF449FEE3C228279BF449FEE3C5136ADBEB9C7EF3C5136ADBEB9C7EF3CE0EB1B3F2DF0F03CE0EB1B3F2DF0F03C70187F3FA118F23C70187F3FA118F23CCD78EF3E1541F33CCD78EF3E1541F33CA06AFBBE8869F43CA06AFBBE8869F43CAD937FBFFB91F53CAD937FBFFB91F53C237816BF6EBAF63C237816BF6EBAF63C47F5B93EE0E2F73C47F5B93EE0E2F73C61F17A3F530BF93C61F17A3F530BF93CE430323FC433FA3CE430323FC433FA3C3A8E69BE365CFB3C3A8E69BE365CFB3C4A4971BFA784FC3C4A4971BFA784FC3C9F584ABF18ADFD3C9F584ABF18ADFD3C620AB53D88D5FE3C620AB53D88D5FE3CE3CC623FF8FDFF3CE3CC623FF8FDFF3C90735E3F3493003D90735E3F3493003D9D4E593D6C27013D9D4E593D6C27013D65C64FBFA3BB013D65C64FBFA3BB013DB31A6EBFDB4F023DB31A6EBFDB4F023D231646BE12E4023D231646BE12E4023D4C97383F4978033D4C97383F4978033DD3FD783F800C043DD3FD783F800C043DD4F0A83EB6A0043DD4F0A83EB6A0043D62B61DBFED34053D62B61DBFED34053D29E57EBF23C9053D29E57EBF23C9053DF674EBBE595D063DF674EBBE595D063DC15AFF3E8FF1063DC15AFF3E8FF1063D73B27F3FC585073D73B27F3FC585073D53A1143FFB19083D53A1143FFB19083D592CBEBE30AE083D592CBEBE30AE083D97617BBF6642093D97617BBF6642093D9D8E30BF9BD6093D9D8E30BF9BD6093D0B5F723ED06A0A3D0B5F723ED06A0A3DB208723F04FF0A3DB208723F04FF0A3D40F3483F39930B3D40F3483F39930B3D0917C7BD6D270C3D0917C7BD6D270C3DA7D763BFA1BB0C3DA7D763BFA1BB0C3D40525DBFD54F0D3D40525DBFD54F0D3D3C1835BD09E40D3D3C1835BD09E40D3D2F17513F3D780E3D2F17513F3D780E3D3C436D3F700C0F3D3C436D3F700C0F3DA92F3D3EA4A00F3DA92F3D3EA4A00F3D5E273ABFD734103D5E273ABFD734103D867478BF09C9103D867478BF09C9103DF3A7A4BE3C5D113DF3A7A4BE3C5D113DBA7D1F3F6FF1113DBA7D1F3F6FF1113DC5AC7E3FA185123DC5AC7E3FA185123D666CE73ED319133D666CE73ED319133DE2A201BF05AE133DE2A201BF05AE133D"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x3727C5AC> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.dynamic_slice %arg0, %arg21, %9, sizes = [1, 8] : (tensor<10x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %13 = stablehlo.reshape %12 : (tensor<1x8xf32>) -> tensor<8xf32> + %14 = stablehlo.dynamic_slice %0, %arg22, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %15 = stablehlo.reshape %14 : (tensor<1x4xf32>) -> tensor<4xf32> + %16 = stablehlo.dynamic_slice %1, %arg22, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %17 = stablehlo.reshape %16 : (tensor<1x4xf32>) -> tensor<4xf32> + %18 = stablehlo.iota dim = 0 : tensor<256xi32> + %19 = stablehlo.broadcast_in_dim %arg23, dims = [] : (tensor) -> tensor<256xi32> + %20 = stablehlo.compare LE, %18, %19, SIGNED : (tensor<256xi32>, tensor<256xi32>) -> tensor<256xi1> + %21 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256xf32> + %22 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256xf32> + %23 = stablehlo.select %20, %21, %22 : tensor<256xi1>, tensor<256xf32> + %24 = stablehlo.multiply %13, %13 : tensor<8xf32> + %25 = stablehlo.reduce(%24 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %26 = stablehlo.divide %25, %7 : tensor + %27 = stablehlo.add %26, %6 : tensor + %28 = stablehlo.rsqrt %27 : tensor + %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor) -> tensor<8xf32> + %30 = stablehlo.multiply %13, %29 : tensor<8xf32> + %31 = stablehlo.multiply %30, %arg5 : tensor<8xf32> + %32 = stablehlo.dot_general %31, %arg10, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %33 = stablehlo.reshape %32 : (tensor<8xf32>) -> tensor<2x4xf32> + %34 = stablehlo.dot_general %31, %arg8, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %35 = stablehlo.reshape %34 : (tensor<4xf32>) -> tensor<1x4xf32> + %36 = stablehlo.dot_general %31, %arg11, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %37 = stablehlo.reshape %36 : (tensor<4xf32>) -> tensor<1x4xf32> + %38 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %39 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %40 = stablehlo.multiply %33, %38 : tensor<2x4xf32> + %41 = stablehlo.slice %33 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %42 = stablehlo.slice %33 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %43 = stablehlo.negate %42 : tensor<2x2xf32> + %44 = stablehlo.concatenate %43, %41, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %45 = stablehlo.multiply %44, %39 : tensor<2x4xf32> + %46 = stablehlo.add %40, %45 : tensor<2x4xf32> + %47 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %48 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %49 = stablehlo.multiply %35, %47 : tensor<1x4xf32> + %50 = stablehlo.slice %35 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %51 = stablehlo.slice %35 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %52 = stablehlo.negate %51 : tensor<1x2xf32> + %53 = stablehlo.concatenate %52, %50, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %54 = stablehlo.multiply %53, %48 : tensor<1x4xf32> + %55 = stablehlo.add %49, %54 : tensor<1x4xf32> + %56 = stablehlo.reshape %55 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %57 = stablehlo.dynamic_update_slice %arg24, %56, %10, %arg23, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %58 = stablehlo.reshape %37 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %59 = stablehlo.dynamic_update_slice %arg25, %58, %10, %arg23, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %60 = stablehlo.slice %57 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %61 = stablehlo.reshape %60 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %62 = stablehlo.slice %59 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %63 = stablehlo.reshape %62 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %64 = stablehlo.reshape %46 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %65 = stablehlo.dot_general %64, %61, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %66 = stablehlo.reshape %65 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %67 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %68 = stablehlo.multiply %66, %67 : tensor<2x256xf32> + %69 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %70 = stablehlo.add %68, %69 : tensor<2x256xf32> + %71 = stablehlo.reduce(%70 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %72 = stablehlo.broadcast_in_dim %71, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %73 = stablehlo.subtract %70, %72 : tensor<2x256xf32> + %74 = stablehlo.exponential %73 : tensor<2x256xf32> + %75 = stablehlo.reduce(%74 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %76 = stablehlo.broadcast_in_dim %75, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %77 = stablehlo.divide %74, %76 : tensor<2x256xf32> + %78 = stablehlo.reshape %77 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %79 = stablehlo.dot_general %78, %63, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %80 = stablehlo.reshape %79 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %81 = stablehlo.reshape %80 : (tensor<2x4xf32>) -> tensor<8xf32> + %82 = stablehlo.dot_general %81, %arg9, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %83 = stablehlo.add %13, %82 : tensor<8xf32> + %84 = stablehlo.multiply %83, %83 : tensor<8xf32> + %85 = stablehlo.reduce(%84 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %86 = stablehlo.divide %85, %7 : tensor + %87 = stablehlo.add %86, %6 : tensor + %88 = stablehlo.rsqrt %87 : tensor + %89 = stablehlo.broadcast_in_dim %88, dims = [] : (tensor) -> tensor<8xf32> + %90 = stablehlo.multiply %83, %89 : tensor<8xf32> + %91 = stablehlo.multiply %90, %arg6 : tensor<8xf32> + %92 = stablehlo.dot_general %91, %arg4, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %93 = stablehlo.dot_general %91, %arg7, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %94 = stablehlo.negate %92 : tensor<16xf32> + %95 = stablehlo.exponential %94 : tensor<16xf32> + %96 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %97 = stablehlo.add %96, %95 : tensor<16xf32> + %98 = stablehlo.divide %96, %97 : tensor<16xf32> + %99 = stablehlo.multiply %92, %98 : tensor<16xf32> + %100 = stablehlo.multiply %99, %93 : tensor<16xf32> + %101 = stablehlo.dot_general %100, %arg3, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %102 = stablehlo.add %83, %101 : tensor<8xf32> + %103 = stablehlo.multiply %102, %102 : tensor<8xf32> + %104 = stablehlo.reduce(%103 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %105 = stablehlo.divide %104, %7 : tensor + %106 = stablehlo.add %105, %6 : tensor + %107 = stablehlo.rsqrt %106 : tensor + %108 = stablehlo.broadcast_in_dim %107, dims = [] : (tensor) -> tensor<8xf32> + %109 = stablehlo.multiply %102, %108 : tensor<8xf32> + %110 = stablehlo.multiply %109, %arg14 : tensor<8xf32> + %111 = stablehlo.dot_general %110, %arg19, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %112 = stablehlo.reshape %111 : (tensor<8xf32>) -> tensor<2x4xf32> + %113 = stablehlo.dot_general %110, %arg17, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %114 = stablehlo.reshape %113 : (tensor<4xf32>) -> tensor<1x4xf32> + %115 = stablehlo.dot_general %110, %arg20, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %116 = stablehlo.reshape %115 : (tensor<4xf32>) -> tensor<1x4xf32> + %117 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %118 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %119 = stablehlo.multiply %112, %117 : tensor<2x4xf32> + %120 = stablehlo.slice %112 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %121 = stablehlo.slice %112 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %122 = stablehlo.negate %121 : tensor<2x2xf32> + %123 = stablehlo.concatenate %122, %120, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %124 = stablehlo.multiply %123, %118 : tensor<2x4xf32> + %125 = stablehlo.add %119, %124 : tensor<2x4xf32> + %126 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %127 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %128 = stablehlo.multiply %114, %126 : tensor<1x4xf32> + %129 = stablehlo.slice %114 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %130 = stablehlo.slice %114 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %131 = stablehlo.negate %130 : tensor<1x2xf32> + %132 = stablehlo.concatenate %131, %129, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %133 = stablehlo.multiply %132, %127 : tensor<1x4xf32> + %134 = stablehlo.add %128, %133 : tensor<1x4xf32> + %135 = stablehlo.reshape %134 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %136 = stablehlo.dynamic_update_slice %57, %135, %11, %arg23, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %137 = stablehlo.reshape %116 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %138 = stablehlo.dynamic_update_slice %59, %137, %11, %arg23, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %139 = stablehlo.slice %136 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %140 = stablehlo.reshape %139 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %141 = stablehlo.slice %138 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %142 = stablehlo.reshape %141 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %143 = stablehlo.reshape %125 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %144 = stablehlo.dot_general %143, %140, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %145 = stablehlo.reshape %144 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %146 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %147 = stablehlo.multiply %145, %146 : tensor<2x256xf32> + %148 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %149 = stablehlo.add %147, %148 : tensor<2x256xf32> + %150 = stablehlo.reduce(%149 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %151 = stablehlo.broadcast_in_dim %150, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %152 = stablehlo.subtract %149, %151 : tensor<2x256xf32> + %153 = stablehlo.exponential %152 : tensor<2x256xf32> + %154 = stablehlo.reduce(%153 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %155 = stablehlo.broadcast_in_dim %154, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %156 = stablehlo.divide %153, %155 : tensor<2x256xf32> + %157 = stablehlo.reshape %156 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %158 = stablehlo.dot_general %157, %142, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %159 = stablehlo.reshape %158 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %160 = stablehlo.reshape %159 : (tensor<2x4xf32>) -> tensor<8xf32> + %161 = stablehlo.dot_general %160, %arg18, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %162 = stablehlo.add %102, %161 : tensor<8xf32> + %163 = stablehlo.multiply %162, %162 : tensor<8xf32> + %164 = stablehlo.reduce(%163 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %165 = stablehlo.divide %164, %7 : tensor + %166 = stablehlo.add %165, %6 : tensor + %167 = stablehlo.rsqrt %166 : tensor + %168 = stablehlo.broadcast_in_dim %167, dims = [] : (tensor) -> tensor<8xf32> + %169 = stablehlo.multiply %162, %168 : tensor<8xf32> + %170 = stablehlo.multiply %169, %arg15 : tensor<8xf32> + %171 = stablehlo.dot_general %170, %arg13, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %172 = stablehlo.dot_general %170, %arg16, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %173 = stablehlo.negate %171 : tensor<16xf32> + %174 = stablehlo.exponential %173 : tensor<16xf32> + %175 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %176 = stablehlo.add %175, %174 : tensor<16xf32> + %177 = stablehlo.divide %175, %176 : tensor<16xf32> + %178 = stablehlo.multiply %171, %177 : tensor<16xf32> + %179 = stablehlo.multiply %178, %172 : tensor<16xf32> + %180 = stablehlo.dot_general %179, %arg12, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %181 = stablehlo.add %162, %180 : tensor<8xf32> + %182 = stablehlo.multiply %181, %181 : tensor<8xf32> + %183 = stablehlo.reduce(%182 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %184 = stablehlo.divide %183, %7 : tensor + %185 = stablehlo.add %184, %6 : tensor + %186 = stablehlo.rsqrt %185 : tensor + %187 = stablehlo.broadcast_in_dim %186, dims = [] : (tensor) -> tensor<8xf32> + %188 = stablehlo.multiply %181, %187 : tensor<8xf32> + %189 = stablehlo.multiply %188, %arg1 : tensor<8xf32> + %190 = stablehlo.dot_general %189, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %191 = stablehlo.constant dense<0> : tensor + %192 = stablehlo.constant dense<0xFF800000> : tensor + %193 = stablehlo.iota dim = 0 : tensor<10xi32> + %194:2 = stablehlo.reduce(%190 init: %192), (%193 init: %191) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %195 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %196 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %197 = stablehlo.or %195, %196 : tensor + %198 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %199 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %200 = stablehlo.and %198, %199 : tensor + %201 = stablehlo.or %197, %200 : tensor + %202 = stablehlo.select %197, %amv_l, %amv_r : tensor, tensor + %203 = stablehlo.select %201, %ami_l, %ami_r : tensor, tensor + stablehlo.return %202, %203 : tensor, tensor + } + return %194#1, %136, %138 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/internlm3/prefill.mlir b/src/lib/mlxcel-xla/assets/internlm3/prefill.mlir new file mode 100644 index 000000000..005476361 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/internlm3/prefill.mlir @@ -0,0 +1,222 @@ +module @prefill { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg4: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg5: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg6: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg7: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg10: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg11: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg12: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg13: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg14: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg15: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg16: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg17: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg18: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg19: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg20: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg21: tensor<256xi32> loc("tokens"), %arg22: tensor<256xi32> loc("positions"), %arg23: tensor loc("real_len")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3F0000803F40510A3F0000803F3311D5BEFFFF7F3F3311D5BEFFFF7F3F26707DBFFEFF7F3F26707DBFFEFF7F3F305527BFFDFF7F3F305527BFFDFF7F3F2C3C913EFCFF7F3F2C3C913EFCFF7F3FB8CD753FFAFF7F3FB8CD753FFAFF7F3FBDFF403FF8FF7F3FBDFF403FF8FF7F3FF6FD14BEF5FF7F3FF6FD14BEF5FF7F3FD53F69BFF2FF7F3FD53F69BFF2FF7F3F64CD56BFEFFF7F3F64CD56BFEFFF7F3F7205913BECFF7F3F7205913BECFF7F3FD006583FE8FF7F3FD006583FE8FF7F3F6F4E683FE4FF7F3F6F4E683FE4FF7F3FD7040C3EDFFF7F3FD7040C3EDFFF7F3FE87A42BFDAFF7F3FE87A42BFDAFF7F3F2C2975BFD5FF7F3F2C2975BFD5FF7F3F36E28CBED0FF7F3F36E28CBED0FF7F3F840A293FCAFF7F3F840A293FCAFF7F3FBF1B7D3FC3FF7F3FBF1B7D3FC3FF7F3F22F0D03EBDFF7F3F22F0D03EBDFF7F3FFC370CBFB6FF7F3FFC370CBFB6FF7F3F6FFD7FBFAFFF7F3F6FFD7FBFAFFF7F3FBF6708BFA7FF7F3FBF6708BFA7FF7F3FFE2DD93E9FFF7F3FFE2DD93E9FFF7F3F78BF7D3F97FF7F3F78BF7D3F97FF7F3F819C253F8FFF7F3F819C253F8FFF7F3F389395BE86FF7F3F389395BE86FF7F3F576D76BF7CFF7F3F576D76BF7CFF7F3FB3803FBF73FF7F3FB3803FBF73FF7F3F18F41D3E69FF7F3F18F41D3E69FF7F3F8E2C6A3F5FFF7F3F8E2C6A3F5FFF7F3FAA8F553F54FF7F3FAA8F553F54FF7F3FB78659BC49FF7F3FB78659BC49FF7F3FE73B59BF3EFF7F3FE73B59BF3EFF7F3F5F5867BF32FF7F3F5F5867BF32FF7F3FEA0803BE27FF7F3FEA0803BE27FF7F3F2DF2433F1AFF7F3F2DF2433F1AFF7F3FB57F743F0EFF7F3FB57F743F0EFF7F3F6C85883E01FF7F3F6C85883E01FF7F3F74BC2ABFF4FE7F3F74BC2ABFF4FE7F3F44C27CBFE6FE7F3F44C27CBFE6FE7F3FE0CACCBED8FE7F3FE0CACCBED8FE7F3FE81B0E3FCAFE7F3FE81B0E3FCAFE7F3FBBF57F3FBBFE7F3FBBF57F3FBBFE7F3F807B063FACFE7F3F807B063FACFE7F3F6D46DDBE9DFE7F3F6D46DDBE9DFE7F3FB3097EBF8DFE7F3FB3097EBF8DFE7F3F80E023BF7DFE7F3F80E023BF7DFE7F3F44E7993E6DFE7F3F44E7993E6DFE7F3F0308773F5DFE7F3F0308773F5DFE7F3FD1FD3D3F4CFE7F3FD1FD3D3F4CFE7F3F0EE726BE3AFE7F3F0EE726BE3AFE7F3F95146BBF29FE7F3F95146BBF29FE7F3FA64D54BF17FE7F3FA64D54BF17FE7F3F2C43B53C04FE7F3F2C43B53C04FE7F3FA26C5A3FF2FD7F3FA26C5A3FF2FD7F3FAC5D663FDFFD7F3FAC5D663FDFFD7F3FB714F43DCCFD7F3FB714F43DCCFD7F3F836545BFB8FD7F3F836545BFB8FD7F3F56D173BFA4FD7F3F56D173BFA4FD7F3FE62584BE90FD7F3FE62584BE90FD7F3FF76A2C3F7BFD7F3FF76A2C3F7BFD7F3FB8637C3F66FD7F3FB8637C3F66FD7F3F83A1C83E51FD7F3F83A1C83E51FD7F3FFAFC0FBF3BFD7F3FFAFC0FBF3BFD7F3FE5E87FBF25FD7F3FE5E87FBF25FD7F3F908C04BF0FFD7F3F908C04BF0FFD7F3F6D5AE13EF8FC7F3F6D5AE13EF8FC7F3FD54E7E3FE1FC7F3FD54E7E3FE1FC7F3F3521223FCAFC7F3F3521223FCAFC7F3F3A389EBEB2FC7F3F3A389EBEB2FC7F3FBC9D77BF9AFC7F3FBC9D77BF9AFC7F3F20773CBF82FC7F3F20773CBF82FC7F3FACD62F3E69FC7F3FACD62F3E69FC7F3FE4F76B3F50FC7F3FE4F76B3F50FC7F3F6107533F37FC7F3F6107533F37FC7F3F5ABFFDBC1DFC7F3F5ABFFDBC1DFC7F3FFC985BBF03FC7F3FFC985BBF03FC7F3F5A5E65BFE9FB7F3F5A5E65BFE9FB7F3FB512E2BDCEFB7F3FB512E2BDCEFB7F3FE4D4463FB3FB7F3FE4D4463FB3FB7F3F141E733F98FB7F3F141E733F98FB7F3F72877F3E7CFB7F3F72877F3E7CFB7F3F05162EBF60FB7F3F05162EBF60FB7F3F1C007CBF44FB7F3F1C007CBF44FB7F3F1F74C4BE27FB7F3F1F74C4BE27FB7F3F28DB113F0AFB7F3F28DB113F0AFB7F3FEDD67F3FEDFA7F3FEDD67F3FEDFA7F3FF69A023FCFFA7F3FF69A023FCFFA7F3FE869E5BEB1FA7F3FE869E5BEB1FA7F3FDE8E7EBF93FA7F3FDE8E7EBF93FA7F3FAA5E20BF74FA7F3FAA5E20BF74FA7F3F0486A23E55FA7F3F0486A23E55FA7F3F7D2E783F36FA7F3F7D2E783F36FA7F3FA8EC3A3F16FA7F3FA8EC3A3F16FA7F3FC2C238BEF6F97F3FC2C238BEF6F97F3F78D66CBFD5F97F3F78D66CBFD5F97F3FE0BC51BFB5F97F3FE0BC51BFB5F97F3F381B233D94F97F3F381B233D94F97F3FEEC05C3F72F97F3FEEC05C3F72F97F3F6E5A643F51F97F3F6E5A643F51F97F3F2A0CD03D2FF97F3F2A0CD03D2FF97F3F494048BF0CF97F3F494048BF0CF97F3FF16572BFE9F87F3FF16572BFE9F87F3FF8BD76BEC6F87F3FF8BD76BEC6F87F3F95BD2F3FA3F87F3F95BD2F3FA3F87F3F71977B3F7FF87F3F71977B3F7FF87F3FCB42C03E5BF87F3FCB42C03E5BF87F3F6AB613BF37F87F3F6AB613BF37F87F3FD4BF7FBF12F87F3FD4BF7FBF12F87F3FBFA600BFEDF77F3FBFA600BFEDF77F3FC874E93EC8F77F3FC874E93EC8F77F3FCBC97E3FA2F77F3FCBC97E3FA2F77F3FE7981E3F7CF77F3FE7981E3F7CF77F3F8BD0A6BE55F77F3F8BD0A6BE55F77F3F44BA78BF2FF77F3F44BA78BF2FF77F3F6F5E39BF07F77F3F6F5E39BF07F77F3F25AB413EE0F67F3F25AB413EE0F67F3F4CB06D3FB8F67F3F4CB06D3FB8F67F3F2A6E503F90F67F3F2A6E503F90F67F3F7E5347BD68F67F3F7E5347BD68F67F3F73E45DBF3FF67F3F73E45DBF3FF67F3FEE5163BF16F67F3FEE5163BF16F67F3F7301BEBDECF57F3F7301BEBDECF57F3FA9A7493FC3F57F3FA9A7493FC3F57F3FF2A8713F99F57F3FF2A8713F99F57F3F8BEF6D3E6EF57F3F8BEF6D3E6EF57F3F9F6131BF43F57F3F9F6131BF43F57F3FBB297BBF18F57F3FBB297BBF18F57F3F9C0DBCBEEDF47F3F9C0DBCBEEDF47F3FB68E153FC1F47F3FB68E153FC1F47F3F99A37F3F95F47F3F99A37F3F95F47F3FE45FFD3E68F47F3FE45FFD3E68F47F3FFA7AEDBE3CF47F3FFA7AEDBE3CF47F3F9DFF7EBF0EF47F3F9DFF7EBF0EF47F3FF6CF1CBFE1F37F3FF6CF1CBFE1F37F3FBA17AB3EB3F37F3FBA17AB3EB3F37F3F0E41793F85F37F3F0E41793F85F37F3F7FCC373F57F37F3F7FCC373F57F37F3FA48F4ABE28F37F3FA48F4ABE28F37F3F5B856EBFF9F27F3F5B856EBFF9F27F3F461B4FBFC9F27F3F461B4FBFC9F27F3FC5876B3D99F27F3FC5876B3D99F27F3F84035F3F69F27F3F84035F3F69F27F3FDF44623F39F27F3FDF44623F39F27F3FEDF2AB3D08F27F3FEDF2AB3D08F27F3FFD0A4BBFD7F17F3FFD0A4BBFD7F17F3F1AE770BFA5F17F3F1AE770BFA5F17F3F591C65BE73F17F3F591C65BE73F17F3F1A02333F41F17F3F1A02333F41F17F3FFCB67A3F0FF17F3FFCB67A3F0FF17F3FA7D4B73EDCF07F3FA7D4B73EDCF07F3F016417BFA9F07F3F016417BFA9F07F3F3E827FBF75F07F3F3E827FBF75F07F3F366DF9BE41F07F3F366DF9BE41F07F3F697CF13E0DF07F3F697CF13E0DF07F3F51307F3FD9EF7F3F51307F3FD9EF7F3FE0031B3FA4EF7F3FE0031B3FA4EF7F3F7A5BAFBE6FEF7F3F7A5BAFBE6FEF7F3FD8C279BF39EF7F3FD8C279BF39EF7F3FDF3636BF03EF7F3FDF3636BF03EF7F3F1470533ECDEE7F3F1470533ECDEE7F3FA2556F3F97EE7F3FA2556F3F97EE7F3F3BC44D3F60EE7F3F3BC44D3F60EE7F3FA9DB87BD29EE7F3FA9DB87BD29EE7F3F1C1E60BFF1ED7F3F1C1E60BFF1ED7F3F463361BFB9ED7F3F463361BFB9ED7F3FF3E099BD81ED7F3FF3E099BD81ED7F3F3F6A4C3F48ED7F3F3F6A4C3F48ED7F3F6D20703F10ED7F3F6D20703F10ED7F3F8E445C3ED6EC7F3F8E445C3ED6EC7F3FFE9E34BF9DEC7F3FFE9E34BF9DEC7F3F353F7ABF63EC7F3F353F7ABF63EC7F3F0298B3BE29EC7F3F0298B3BE29EC7F3F4336193FEEEB7F3F4336193FEEEB7F3FC35B7F3FB3EB7F3FC35B7F3FB3EB7F3F8875F53E78EB7F3F8875F53E78EB7F3FFF78F5BE3DEB7F3FFF78F5BE3DEB7F3FE75B7FBF01EB7F3FE75B7FBF01EB7F3FAE3419BFC4EA7F3FAE3419BFC4EA7F3FB69BB33E88EA7F3FB69BB33E88EA7F3F9F3F7A3F4BEA7F3F9F3F7A3F4BEA7F3F979D343F0EEA7F3F979D343F0EEA7F3F464C5CBED0E97F3F464C5CBED0E97F3F1C2170BF92E97F3F1C2170BF92E97F3F0F694CBF54E97F3F0F694CBF54E97F3FB6F0993D16E97F3FB6F0993D16E97F3F3634613FD7E87F3F3634613FD7E87F3F281D603F97E87F3F281D603F97E87F3FE4CB873D58E87F3FE4CB873D58E87F3F67C54DBF18E87F3F67C54DBF18E87F3FEF546FBFD8E77F3FEF546FBFD8E77F3F586853BE97E77F3F586853BE97E77F3F4238363F56E77F3F4238363F56E77F3F69C2793F15E77F3F69C2793F15E77F3FC357AF3ED3E67F3FC357AF3ED3E67F3F72051BBF91E67F3F72051BBF91E67F3F29307FBF4FE67F3F29307FBF4FE67F3FED78F1BE0CE67F3FED78F1BE0CE67F3FAA70F93ECAE57F3FAA70F93ECAE57F3F5D827F3F86E57F3F5D827F3F86E57F3F6962173F43E57F3F6962173F43E57F3F57D8B7BEFFE47F3F57D8B7BEFFE47F3F62B77ABFBAE47F3F62B77ABFBAE47F3FB00033BF76E47F3FB00033BF76E47F3F0D24653E31E47F3F0D24653E31E47F3FC5E7703FECE37F3FC5E7703FECE37F3FC9094B3FA6E37F3FC9094B3FA6E37F3FAC02ACBD60E37F3FAC02ACBD60E37F3FCB4562BF1AE37F3FCB4562BF1AE37F3F8C025FBFD3E27F3F8C025FBFD3E27F3F36686BBD8CE27F3F36686BBD8CE27F3F6F1C4F3F45E27F3F6F1C4F3F45E27F3FA4846E3FFDE17F3FA4846E3FFDE17F3FE5874A3EB5E17F3FE5874A3EB5E17F3FDFCD37BF6DE17F3FDFCD37BF6DE17F3F9A4079BF24E17F3F9A4079BF24E17F3F0014ABBEDBE07F3F0014ABBEDBE07F3F86D11C3F92E07F3F86D11C3F92E07F3F70FF7E3F48E07F3F70FF7E3F48E07F3F7A77ED3EFFDF7F3F7A77ED3EFFDF7F3F5363FDBEB4DF7F3F5363FDBEB4DF7F3FB4A37FBF6ADF7F3FB4A37FBF6ADF7F3F1B8D15BF1FDF7F3F1B8D15BF1FDF7F3F4911BC3ED3DE7F3F4911BC3ED3DE7F3F1D2A7B3F88DE7F3F1D2A7B3F88DE7F3F3260313F3CDE7F3F3260313F3CDE7F3F3BF76DBEEFDD7F3F3BF76DBEEFDD7F3F99A971BFA3DD7F3F99A971BFA3DD7F3F71A649BF56DD7F3F71A649BF56DD7F3F3011BE3D08DD7F3F3011BE3D08DD7F3FD752633FBBDC7F3FD752633FBBDC7F3F76E35D3F6DDC7F3F76E35D3F6DDC7F3FEC33473D1EDC7F3FEC33473D1EDC7F3F506F50BFD0DB7F3F506F50BFD0DB7F3F90AF6DBF81DB7F3F90AF6DBF81DB7F3F62A341BE31DB7F3F62A341BE31DB7F3FCC5F393FE2DA7F3FCC5F393FE2DA7F3FCCB9783F92DA7F3FCCB9783F92DA7F3FCFCCA63E41DA7F3FCFCCA63E41DA7F3F749A1EBFF1D97F3F749A1EBFF1D97F3F9AC97EBFA0D97F3F9AC97EBFA0D97F3F4471E9BE4ED97F3F4471E9BE4ED97F3F74A8003FFDD87F3F74A8003FFDD87F3FEABF7F3FAAD87F3FEABF7F3FAAD87F3FCDB4133F58D87F3FCDB4133F58D87F3F7546C0BE05D87F3F7546C0BE05D87F3FCF977BBFB2D77F3FCF977BBFB2D77F3F25BC2FBF5FD77F3F25BC2FBF5FD77F3FA3C5763E0BD77F3FA3C5763E0BD77F3F9466723FB7D67F3F9466723FB7D67F3F0E3F483F63D67F3F0E3F483F63D67F3FE31BD0BD0ED67F3FE31BD0BD0ED67F3F535B64BFB9D57F3F535B64BFB9D57F3FEEBF5CBF64D57F3FEEBF5CBF64D57F3F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F824A1439A46A573F824A1439B7C7683F814A9439B7C7683F814A9439C381103EC26FDE39C381103EC26FDE39CFBD41BF814A143ACFBD41BF814A143A107C75BF215D393A107C75BF215D393A8C0F8FBEC16F5E3A8C0F8FBEC16F5E3A4630283F30C1813A4630283F30C1813A95467D3F804A943A95467D3F804A943A3201D33ECFD3A63A3201D33ECFD3A63AF8440BBF1E5DB93AF8440BBF1E5DB93A5CFF7FBF6DE6CB3A5CFF7FBF6DE6CB3AD85C09BFBB6FDE3AD85C09BFBB6FDE3A2220D73E0AF9F03A2220D73E0AF9F03A72987D3F2CC1013B72987D3F2CC1013B4479263FD3050B3B4479263FD3050B3B106893BE794A143B106893BE794A143B251E76BF208F1D3B251E76BF208F1D3BB34040BFC6D3263BB34040BFC6D3263B6979193E6C18303B6979193E6C18303BC8B6693F125D393BC8B6693F125D393B102F563FB7A1423B102F563FB7A1423B150511BC5DE64B3B150511BC5DE64B3BE7A158BF022B553BE7A158BF022B553BFCD367BFA66F5E3BFCD367BFA66F5E3B388707BE4BB4673B388707BE4BB4673B0837433FEFF8703B0837433FEFF8703B0ED5743F933D7A3B0ED5743F933D7A3B2AB48A3E1BC1813B2AB48A3E1BC1813BE9E329BF6D63863BE9E329BF6D63863BA4EF7CBFBE058B3BA4EF7CBFBE058B3B06DECEBE0FA88F3B06DECEBE0FA88F3B4D2A0D3F604A943B4D2A0D3F604A943B39FA7F3FB1EC983B39FA7F3FB1EC983BF771073F028F9D3BF771073F028F9D3BC23ADBBE5231A23BC23ADBBE5231A23B38E57DBFA3D3A63B38E57DBFA3D3A63BEABE24BFF375AB3BEABE24BFF375AB3B9FBD973E4218B03B9FBD973E4218B03B4BBB763F92BAB43B4BBB763F92BAB43BBCBF3E3FE15CB93BBCBF3E3FE15CB93BFB6D22BE30FFBD3BFB6D22BE30FFBD3B28A16ABF7FA1C23B28A16ABF7FA1C23B31EF54BFCE43C73B31EF54BFCE43C73BA103913C1CE6CB3BA103913C1CE6CB3BD0D4593F6A88D03BD0D4593F6A88D03B9ADB663FB82AD53B9ADB663FB82AD53BE813FD3D05CDD93BE813FD3D05CDD93B56AC44BF526FDE3B56AC44BF526FDE3B222974BF9F11E33B222974BF9F11E33BFF5586BEECB3E73BFF5586BEECB3E73B24942B3F3856EC3B24942B3F3856EC3BA0937C3F84F8F03BA0937C3F84F8F03BB4B6CA3ED09AF53BB4B6CA3ED09AF53BCD0C0FBF1B3DFA3BCD0C0FBF1B3DFA3BF4EF7FBF66DFFE3BF4EF7FBF66DFFE3B5E8405BFD9C0013C5E8405BFD9C0013CFD50DF3EFE11043CFD50DF3EFE11043CE72C7E3F2363063CE72C7E3F2363063C4301233F48B4083C4301233F48B4083C23109CBE6C050B3C23109CBE6C050B3C7F5377BF91560D3C7F5377BF91560D3CF23A3DBFB5A70F3CF23A3DBFB5A70F3C4B5F2B3ED9F8113C4B5F2B3ED9F8113CD4866B3FFD49143CD4866B3FFD49143C0BAB533F219B163C0BAB533F219B163CCE81D9BC44EC183CCE81D9BC44EC183C5B035BBF673D1B3C5B035BBF673D1B3C97DE65BF8B8E1D3C97DE65BF8B8E1D3C4D14EBBDADDF1F3C4D14EBBDADDF1F3CB31D463FD030223CB31D463FD030223C5178733FF381243C5178733FF381243C23F5813E15D3263C23F5813E15D3263CED402DBF3724293CED402DBF3724293C8C327CBF59752B3C8C327CBF59752B3C518BC6BE7AC62D3C518BC6BE7AC62D3C6EEC103F9C17303C6EEC103F9C17303C8EE07F3FBD68323C8EE07F3FBD68323C1894033FDEB9343C1894033FDEB9343CBC62E3BEFE0A373CBC62E3BEFE0A373C7D6F7EBF1F5C393C7D6F7EBF1F5C393C574021BF3FAD3B3C574021BF3FAD3B3C865FA03E5FFE3D3C865FA03E5FFE3D3CBCE6773F7F4F403CBCE6773F7F4F403C5CB23B3F9EA0423C5CB23B3F9EA0423C2B4D34BEBDF1443C2B4D34BEBDF1443CC6676CBFDC42473CC6676CBFDC42473CA76252BFFB93493CA76252BFFB93493CD0FD103D19E54B3CD0FD103D19E54B3C822D5C3F37364E3C822D5C3F37364E3CF7DC643F5587503CF7DC643F5587503CFB0FD93D73D8523CFB0FD93D73D8523C178B47BF9029553C178B47BF9029553C9EC272BFAD7A573C9EC272BFAD7A573C56237BBECACB593C56237BBECACB593C3DEA2E3FE61C5C3C3DEA2E3FE61C5C3C68CC7B3F036E5E3C68CC7B3F036E5E3CF25BC23E1EBF603CF25BC23E1EBF603C28C912BF3A10633C28C912BF3A10633C05CC7FBF5561653C05CC7FBF5561653C2EA101BF70B2673C2EA101BF70B2673CED6FE73E8B036A3CED6FE73E8B036A3CF8AC7E3FA5546C3CF8AC7E3FA5546C3C2F7C1F3FBFA56E3C2F7C1F3FBFA56E3CB1ABA4BED9F6703CB1ABA4BED9F6703C007578BFF347733C007578BFF347733C03263ABF0C99753C03263ABF0C99753C6D373D3E25EA773C6D373D3E25EA773CFA436D3F3D3B7A3CFA436D3F3D3B7A3C0B16513F558C7C3C0B16513F558C7C3CCF3735BD6DDD7E3CCF3735BD6DDD7E3C3E535DBF4297803C3E535DBF4297803CC0D663BFCEBF813CC0D663BFCEBF813C4F07C7BD59E8823C4F07C7BD59E8823C7AF4483FE410843C7AF4483FE410843C0D08723F7039853C0D08723F7039853C5D57723EFA61863C5D57723EFA61863C0B9030BF858A873C0B9030BF858A873C38617BBF10B3883C38617BBF10B3883CAE28BEBE9ADB893CAE28BEBE9ADB893CEFA2143F24048B3CEFA2143F24048B3C5BB27F3FAE2C8C3C5BB27F3FAE2C8C3C5557FF3E38558D3C5557FF3E38558D3C7878EBBEC27D8E3C7878EBBEC27D8E3C58E57EBF4BA68F3C58E57EBF4BA68F3CD4B41DBFD4CE903CD4B41DBFD4CE903C8FF4A83E5DF7913C8FF4A83E5DF7913C49FE783FE61F933C49FE783FE61F933CED95383F6F48943CED95383F6F48943CE41D46BEF770953CE41D46BEF770953C6C1B6EBF8099963C6C1B6EBF8099963C3DC54FBF08C2973C3DC54FBF08C2973C2D6E593D90EA983C2D6E593D90EA983C8A745E3F17139A3C8A745E3F17139A3CF8CB623F9F3B9B3CF8CB623F9F3B9B3CA4FAB43D26649C3CA4FAB43D26649C3CD5594ABFAD8C9D3CD5594ABFAD8C9D3CA14871BF34B59E3CA14871BF34B59E3C888669BEBBDD9F3C888669BEBBDD9F3C4F32323F4106A13C4F32323F4106A13CFCF07A3FC72EA23CFCF07A3FC72EA23C99F1B93E4D57A33C99F1B93E4D57A33CBC7916BFD37FA43CBC7916BFD37FA43C90937FBF59A8A53C90937FBF59A8A53C2F67FBBEDED0A63C2F67FBBEDED0A63C4B7CEF3E63F9A73C4B7CEF3E63F9A73C9B187F3FE821A93C9B187F3FE821A93C4FEA1B3F6D4AAA3C4FEA1B3F6D4AAA3C093AADBEF172AB3C093AADBEF172AB3C938279BF769BAC3C938279BF769BAC3C240237BFFAC3AD3C240237BFFAC3AD3C61004F3E7EECAE3C61004F3E7EECAE3C18EE6E3F0115B03C18EE6E3F0115B03C45704E3F843DB13C45704E3F843DB13C2EA07DBD0866B23C2EA07DBD0866B23C60915FBF8B8EB33C60915FBF8B8EB33CA3BC61BF0DB7B43CA3BC61BF0DB7B43C59EAA2BD90DFB53C59EAA2BD90DFB53C21BB4B3F1208B73C21BB4B3F1208B73C5D84703F9430B83C5D84703F9430B83C04B1603E1559B93C04B1603E1559B93CFFD033BF9781BA3CFFD033BF9781BA3CB97B7ABF18AABB3CB97B7ABF18AABB3CC9B6B5BE99D2BC3CC9B6B5BE99D2BC3C844D183F1AFBBD3C844D183F1AFBBD3CA46F7F3F9A23BF3CA46F7F3F9A23BF3CFE71F73E1A4CC03CFE71F73E1A4CC03C507BF3BE9A74C13C507BF3BE9A74C13CC0467FBF1A9DC23CC0467FBF1A9DC23CAA1C1ABF9AC5C33CAA1C1ABF9AC5C33C0A7CB13E19EEC43C0A7CB13E19EEC43CDC017A3F9816C63CDC017A3F9816C63CAF6A353F163FC73CAF6A353F163FC73CB7DE57BE9567C83CB7DE57BE9567C83CF9BB6FBF1390C93CF9BB6FBF1390C93C28174DBF91B8CA3C28174DBF91B8CA3C8CE6903D0EE1CB3C8CE6903D0EE1CB3CB9A9603F8C09CD3CB9A9603F8C09CD3CC7A8603F0932CE3CC7A8603F0932CE3CC9D6903D865ACF3CC9D6903D865ACF3C57184DBF0283D03C57184DBF0283D03C47BB6FBF7FABD13C47BB6FBF7FABD13CFED657BEFBD3D23CFED657BEFBD3D23C146C353F76FCD33C146C353F76FCD33C6F017A3FF224D53C6F017A3FF224D53C5578B13E6D4DD63C5578B13E6D4DD63C3E1E1ABFE875D73C3E1E1ABFE875D73C9A467FBF629ED83C9A467FBF629ED83CD777F3BEDDC6D93CD777F3BEDDC6D93C7375F73E57EFDA3C7375F73E57EFDA3CC66F7F3FD117DC3CC66F7F3FD117DC3CED4B183F4A40DD3CED4B183F4A40DD3C7BBAB5BEC368DE3C7BBAB5BEC368DE3C217C7ABF3C91DF3C217C7ABF3C91DF3C97CF33BFB5B9E03C97CF33BFB5B9E03CB9B8603E2DE2E13CB9B8603E2DE2E13C0B85703FA50AE33C0B85703FA50AE33CEFB94B3F1D33E43CEFB94B3F1D33E43C1AFAA2BD945BE53C1AFAA2BD945BE53C92BD61BF0B84E63C92BD61BF0B84E63C69905FBF82ACE73C69905FBF82ACE73CA2807DBDF8D4E83CA2807DBDF8D4E83C70714E3F6FFDE93C70714E3F6FFDE93C62ED6E3FE425EB3C62ED6E3FE425EB3CA4F84E3E5A4EEC3CA4F84E3E5A4EEC3C860337BFCF76ED3C860337BFCF76ED3C228279BF449FEE3C228279BF449FEE3C5136ADBEB9C7EF3C5136ADBEB9C7EF3CE0EB1B3F2DF0F03CE0EB1B3F2DF0F03C70187F3FA118F23C70187F3FA118F23CCD78EF3E1541F33CCD78EF3E1541F33CA06AFBBE8869F43CA06AFBBE8869F43CAD937FBFFB91F53CAD937FBFFB91F53C237816BF6EBAF63C237816BF6EBAF63C47F5B93EE0E2F73C47F5B93EE0E2F73C61F17A3F530BF93C61F17A3F530BF93CE430323FC433FA3CE430323FC433FA3C3A8E69BE365CFB3C3A8E69BE365CFB3C4A4971BFA784FC3C4A4971BFA784FC3C9F584ABF18ADFD3C9F584ABF18ADFD3C620AB53D88D5FE3C620AB53D88D5FE3CE3CC623FF8FDFF3CE3CC623FF8FDFF3C90735E3F3493003D90735E3F3493003D9D4E593D6C27013D9D4E593D6C27013D65C64FBFA3BB013D65C64FBFA3BB013DB31A6EBFDB4F023DB31A6EBFDB4F023D231646BE12E4023D231646BE12E4023D4C97383F4978033D4C97383F4978033DD3FD783F800C043DD3FD783F800C043DD4F0A83EB6A0043DD4F0A83EB6A0043D62B61DBFED34053D62B61DBFED34053D29E57EBF23C9053D29E57EBF23C9053DF674EBBE595D063DF674EBBE595D063DC15AFF3E8FF1063DC15AFF3E8FF1063D73B27F3FC585073D73B27F3FC585073D53A1143FFB19083D53A1143FFB19083D592CBEBE30AE083D592CBEBE30AE083D97617BBF6642093D97617BBF6642093D9D8E30BF9BD6093D9D8E30BF9BD6093D0B5F723ED06A0A3D0B5F723ED06A0A3DB208723F04FF0A3DB208723F04FF0A3D40F3483F39930B3D40F3483F39930B3D0917C7BD6D270C3D0917C7BD6D270C3DA7D763BFA1BB0C3DA7D763BFA1BB0C3D40525DBFD54F0D3D40525DBFD54F0D3D3C1835BD09E40D3D3C1835BD09E40D3D2F17513F3D780E3D2F17513F3D780E3D3C436D3F700C0F3D3C436D3F700C0F3DA92F3D3EA4A00F3DA92F3D3EA4A00F3D5E273ABFD734103D5E273ABFD734103D867478BF09C9103D867478BF09C9103DF3A7A4BE3C5D113DF3A7A4BE3C5D113DBA7D1F3F6FF1113DBA7D1F3F6FF1113DC5AC7E3FA185123DC5AC7E3FA185123D666CE73ED319133D666CE73ED319133DE2A201BF05AE133DE2A201BF05AE133D"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x3727C5AC> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.reshape %arg21 : (tensor<256xi32>) -> tensor<256x1xi32> + %13 = "stablehlo.gather"(%arg0, %12) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<10x8xf32>, tensor<256x1xi32>) -> tensor<256x8xf32> + %14 = stablehlo.reshape %arg22 : (tensor<256xi32>) -> tensor<256x1xi32> + %15 = "stablehlo.gather"(%0, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %16 = "stablehlo.gather"(%1, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %17 = stablehlo.iota dim = 0 : tensor<256xi32> + %18 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<256xi32>) -> tensor<256x256xi32> + %19 = stablehlo.iota dim = 0 : tensor<256xi32> + %20 = stablehlo.broadcast_in_dim %19, dims = [1] : (tensor<256xi32>) -> tensor<256x256xi32> + %21 = stablehlo.compare LE, %20, %18, SIGNED : (tensor<256x256xi32>, tensor<256x256xi32>) -> tensor<256x256xi1> + %22 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x256xf32> + %23 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256x256xf32> + %24 = stablehlo.select %21, %22, %23 : tensor<256x256xi1>, tensor<256x256xf32> + %25 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %26 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %27 = stablehlo.multiply %13, %13 : tensor<256x8xf32> + %28 = stablehlo.reduce(%27 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %29 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %30 = stablehlo.divide %28, %29 : tensor<256xf32> + %31 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %32 = stablehlo.add %30, %31 : tensor<256xf32> + %33 = stablehlo.rsqrt %32 : tensor<256xf32> + %34 = stablehlo.broadcast_in_dim %33, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %35 = stablehlo.multiply %13, %34 : tensor<256x8xf32> + %36 = stablehlo.broadcast_in_dim %arg5, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %37 = stablehlo.multiply %35, %36 : tensor<256x8xf32> + %38 = stablehlo.dot_general %37, %arg10, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %39 = stablehlo.reshape %38 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %40 = stablehlo.dot_general %37, %arg8, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %41 = stablehlo.reshape %40 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %42 = stablehlo.dot_general %37, %arg11, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %43 = stablehlo.reshape %42 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %44 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %45 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %46 = stablehlo.multiply %39, %44 : tensor<256x2x4xf32> + %47 = stablehlo.slice %39 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %48 = stablehlo.slice %39 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %49 = stablehlo.negate %48 : tensor<256x2x2xf32> + %50 = stablehlo.concatenate %49, %47, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %51 = stablehlo.multiply %50, %45 : tensor<256x2x4xf32> + %52 = stablehlo.add %46, %51 : tensor<256x2x4xf32> + %53 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %54 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %55 = stablehlo.multiply %41, %53 : tensor<256x1x4xf32> + %56 = stablehlo.slice %41 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %57 = stablehlo.slice %41 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %58 = stablehlo.negate %57 : tensor<256x1x2xf32> + %59 = stablehlo.concatenate %58, %56, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %60 = stablehlo.multiply %59, %54 : tensor<256x1x4xf32> + %61 = stablehlo.add %55, %60 : tensor<256x1x4xf32> + %62 = stablehlo.reshape %61 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %63 = stablehlo.dynamic_update_slice %25, %62, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %64 = stablehlo.reshape %43 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %65 = stablehlo.dynamic_update_slice %26, %64, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %66 = stablehlo.reshape %52 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %67 = stablehlo.dot_general %66, %61, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %68 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %69 = stablehlo.multiply %67, %68 : tensor<1x256x2x256xf32> + %70 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %71 = stablehlo.add %69, %70 : tensor<1x256x2x256xf32> + %72 = stablehlo.reduce(%71 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %73 = stablehlo.broadcast_in_dim %72, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %74 = stablehlo.subtract %71, %73 : tensor<1x256x2x256xf32> + %75 = stablehlo.exponential %74 : tensor<1x256x2x256xf32> + %76 = stablehlo.reduce(%75 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %77 = stablehlo.broadcast_in_dim %76, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %78 = stablehlo.divide %75, %77 : tensor<1x256x2x256xf32> + %79 = stablehlo.dot_general %78, %43, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %80 = stablehlo.transpose %79, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %81 = stablehlo.reshape %80 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %82 = stablehlo.dot_general %81, %arg9, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %83 = stablehlo.add %13, %82 : tensor<256x8xf32> + %84 = stablehlo.multiply %83, %83 : tensor<256x8xf32> + %85 = stablehlo.reduce(%84 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %86 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %87 = stablehlo.divide %85, %86 : tensor<256xf32> + %88 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %89 = stablehlo.add %87, %88 : tensor<256xf32> + %90 = stablehlo.rsqrt %89 : tensor<256xf32> + %91 = stablehlo.broadcast_in_dim %90, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %92 = stablehlo.multiply %83, %91 : tensor<256x8xf32> + %93 = stablehlo.broadcast_in_dim %arg6, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %94 = stablehlo.multiply %92, %93 : tensor<256x8xf32> + %95 = stablehlo.dot_general %94, %arg4, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %96 = stablehlo.dot_general %94, %arg7, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %97 = stablehlo.negate %95 : tensor<256x16xf32> + %98 = stablehlo.exponential %97 : tensor<256x16xf32> + %99 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %100 = stablehlo.add %99, %98 : tensor<256x16xf32> + %101 = stablehlo.divide %99, %100 : tensor<256x16xf32> + %102 = stablehlo.multiply %95, %101 : tensor<256x16xf32> + %103 = stablehlo.multiply %102, %96 : tensor<256x16xf32> + %104 = stablehlo.dot_general %103, %arg3, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %105 = stablehlo.add %83, %104 : tensor<256x8xf32> + %106 = stablehlo.multiply %105, %105 : tensor<256x8xf32> + %107 = stablehlo.reduce(%106 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %108 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %109 = stablehlo.divide %107, %108 : tensor<256xf32> + %110 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %111 = stablehlo.add %109, %110 : tensor<256xf32> + %112 = stablehlo.rsqrt %111 : tensor<256xf32> + %113 = stablehlo.broadcast_in_dim %112, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %114 = stablehlo.multiply %105, %113 : tensor<256x8xf32> + %115 = stablehlo.broadcast_in_dim %arg14, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %116 = stablehlo.multiply %114, %115 : tensor<256x8xf32> + %117 = stablehlo.dot_general %116, %arg19, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %118 = stablehlo.reshape %117 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %119 = stablehlo.dot_general %116, %arg17, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %120 = stablehlo.reshape %119 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %121 = stablehlo.dot_general %116, %arg20, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %122 = stablehlo.reshape %121 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %123 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %124 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %125 = stablehlo.multiply %118, %123 : tensor<256x2x4xf32> + %126 = stablehlo.slice %118 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %127 = stablehlo.slice %118 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %128 = stablehlo.negate %127 : tensor<256x2x2xf32> + %129 = stablehlo.concatenate %128, %126, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %130 = stablehlo.multiply %129, %124 : tensor<256x2x4xf32> + %131 = stablehlo.add %125, %130 : tensor<256x2x4xf32> + %132 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %133 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %134 = stablehlo.multiply %120, %132 : tensor<256x1x4xf32> + %135 = stablehlo.slice %120 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %136 = stablehlo.slice %120 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %137 = stablehlo.negate %136 : tensor<256x1x2xf32> + %138 = stablehlo.concatenate %137, %135, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %139 = stablehlo.multiply %138, %133 : tensor<256x1x4xf32> + %140 = stablehlo.add %134, %139 : tensor<256x1x4xf32> + %141 = stablehlo.reshape %140 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %142 = stablehlo.dynamic_update_slice %63, %141, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %143 = stablehlo.reshape %122 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %144 = stablehlo.dynamic_update_slice %65, %143, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %145 = stablehlo.reshape %131 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %146 = stablehlo.dot_general %145, %140, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %147 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %148 = stablehlo.multiply %146, %147 : tensor<1x256x2x256xf32> + %149 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %150 = stablehlo.add %148, %149 : tensor<1x256x2x256xf32> + %151 = stablehlo.reduce(%150 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %152 = stablehlo.broadcast_in_dim %151, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %153 = stablehlo.subtract %150, %152 : tensor<1x256x2x256xf32> + %154 = stablehlo.exponential %153 : tensor<1x256x2x256xf32> + %155 = stablehlo.reduce(%154 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %156 = stablehlo.broadcast_in_dim %155, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %157 = stablehlo.divide %154, %156 : tensor<1x256x2x256xf32> + %158 = stablehlo.dot_general %157, %122, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %159 = stablehlo.transpose %158, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %160 = stablehlo.reshape %159 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %161 = stablehlo.dot_general %160, %arg18, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %162 = stablehlo.add %105, %161 : tensor<256x8xf32> + %163 = stablehlo.multiply %162, %162 : tensor<256x8xf32> + %164 = stablehlo.reduce(%163 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %165 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %166 = stablehlo.divide %164, %165 : tensor<256xf32> + %167 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %168 = stablehlo.add %166, %167 : tensor<256xf32> + %169 = stablehlo.rsqrt %168 : tensor<256xf32> + %170 = stablehlo.broadcast_in_dim %169, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %171 = stablehlo.multiply %162, %170 : tensor<256x8xf32> + %172 = stablehlo.broadcast_in_dim %arg15, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %173 = stablehlo.multiply %171, %172 : tensor<256x8xf32> + %174 = stablehlo.dot_general %173, %arg13, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %175 = stablehlo.dot_general %173, %arg16, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %176 = stablehlo.negate %174 : tensor<256x16xf32> + %177 = stablehlo.exponential %176 : tensor<256x16xf32> + %178 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %179 = stablehlo.add %178, %177 : tensor<256x16xf32> + %180 = stablehlo.divide %178, %179 : tensor<256x16xf32> + %181 = stablehlo.multiply %174, %180 : tensor<256x16xf32> + %182 = stablehlo.multiply %181, %175 : tensor<256x16xf32> + %183 = stablehlo.dot_general %182, %arg12, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %184 = stablehlo.add %162, %183 : tensor<256x8xf32> + %185 = stablehlo.multiply %184, %184 : tensor<256x8xf32> + %186 = stablehlo.reduce(%185 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %187 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %188 = stablehlo.divide %186, %187 : tensor<256xf32> + %189 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %190 = stablehlo.add %188, %189 : tensor<256xf32> + %191 = stablehlo.rsqrt %190 : tensor<256xf32> + %192 = stablehlo.broadcast_in_dim %191, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %193 = stablehlo.multiply %184, %192 : tensor<256x8xf32> + %194 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %195 = stablehlo.multiply %193, %194 : tensor<256x8xf32> + %196 = stablehlo.constant dense<1> : tensor + %197 = stablehlo.subtract %arg23, %196 : tensor + %198 = stablehlo.dynamic_slice %195, %197, %9, sizes = [1, 8] : (tensor<256x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %199 = stablehlo.reshape %198 : (tensor<1x8xf32>) -> tensor<8xf32> + %200 = stablehlo.dot_general %199, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %201 = stablehlo.constant dense<0> : tensor + %202 = stablehlo.constant dense<0xFF800000> : tensor + %203 = stablehlo.iota dim = 0 : tensor<10xi32> + %204:2 = stablehlo.reduce(%200 init: %202), (%203 init: %201) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %205 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %206 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %207 = stablehlo.or %205, %206 : tensor + %208 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %209 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %210 = stablehlo.and %208, %209 : tensor + %211 = stablehlo.or %207, %210 : tensor + %212 = stablehlo.select %207, %amv_l, %amv_r : tensor, tensor + %213 = stablehlo.select %211, %ami_l, %ami_r : tensor, tensor + stablehlo.return %212, %213 : tensor, tensor + } + return %204#1, %142, %144 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/mimo/README.md b/src/lib/mlxcel-xla/assets/mimo/README.md new file mode 100644 index 000000000..7a2f39425 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/mimo/README.md @@ -0,0 +1,18 @@ +# MiMo config fixture (#499 dense arch pack) + +`config.json` is a **small synthetic** MiMo config (real switches, tiny dims) +registered as a byte-exact structural fixture in `validation.rs`; the real MiMo-7B +`config.json` parsing is covered by the `config::tests` parse assertion. + +MiMo (`model_type` `mimo`) IS the **Qwen2** architecture: `MiMoForCausalLM` subclasses +`Qwen2ForCausalLM` and reuses `Qwen2Attention` / `Qwen2MLP` / `Qwen2RMSNorm` verbatim, +so it has a **q/k/v projection bias** and **untied** embeddings with plain RoPE +(half-split, `head_dim^-0.5` scaling). The extra multi-token-prediction (`MiMoMTPLayers`, +`num_nextn_predict_layers`) heads are not part of the base transformer and are not +loaded. Its config `sliding_window` is served globally (as for Qwen2), so it parses +to `sliding_window = None`. + +Validation: `validation::dense_pack_families_reuse_proven_graphs` asserts it emits +byte-for-byte identical StableHLO to an untied Qwen2 reference; since MiMo's decoder +is Qwen2 verbatim, the existing Qwen2 token-exact gate covers its forward. The full +gate runs post-merge on the real checkpoint via `scripts/xla/validate_arch.sh`. diff --git a/src/lib/mlxcel-xla/assets/mimo/config.json b/src/lib/mlxcel-xla/assets/mimo/config.json new file mode 100644 index 000000000..5c7fb4127 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/mimo/config.json @@ -0,0 +1,17 @@ +{ + "model_type": "mimo", + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 16, + "num_hidden_layers": 2, + "rms_norm_eps": 1e-05, + "rope_theta": 640000.0, + "vocab_size": 10, + "tie_word_embeddings": false, + "attention_bias": true, + "sliding_window": 32768, + "use_sliding_window": true, + "hidden_act": "silu" +} diff --git a/src/lib/mlxcel-xla/assets/mimo/decode.mlir b/src/lib/mlxcel-xla/assets/mimo/decode.mlir new file mode 100644 index 000000000..db230ad8e --- /dev/null +++ b/src/lib/mlxcel-xla/assets/mimo/decode.mlir @@ -0,0 +1,218 @@ +module @decode_step { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg4: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg5: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg6: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg7: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg10: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg11: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg12: tensor<4xf32> loc("params['layers'][0]['bk']"), %arg13: tensor<8xf32> loc("params['layers'][0]['bq']"), %arg14: tensor<4xf32> loc("params['layers'][0]['bv']"), %arg15: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg16: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg17: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg18: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg19: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg20: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg21: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg22: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg23: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg24: tensor<4xf32> loc("params['layers'][1]['bk']"), %arg25: tensor<8xf32> loc("params['layers'][1]['bq']"), %arg26: tensor<4xf32> loc("params['layers'][1]['bv']"), %arg27: tensor loc("token"), %arg28: tensor loc("pos"), %arg29: tensor loc("cache_len"), %arg30: tensor<2x256x1x4xf32> loc("kcache"), %arg31: tensor<2x256x1x4xf32> loc("vcache")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3FF3FF7F3F40510A3FF3FF7F3F3311D5BECCFF7F3F3311D5BECCFF7F3F26707DBF8AFF7F3F26707DBF8AFF7F3F305527BF2EFF7F3F305527BF2EFF7F3F2C3C913EB8FE7F3F2C3C913EB8FE7F3FB8CD753F28FE7F3FB8CD753F28FE7F3FBDFF403F7EFD7F3FBDFF403F7EFD7F3FF6FD14BEB9FC7F3FF6FD14BEB9FC7F3FD53F69BFDAFB7F3FD53F69BFDAFB7F3F64CD56BFE1FA7F3F64CD56BFE1FA7F3F7205913BCEF97F3F7205913BCEF97F3FD006583FA1F87F3FD006583FA1F87F3F6F4E683F59F77F3F6F4E683F59F77F3FD7040C3EF7F57F3FD7040C3EF7F57F3FE87A42BF7BF47F3FE87A42BF7BF47F3F2C2975BFE5F27F3F2C2975BFE5F27F3F36E28CBE34F17F3F36E28CBE34F17F3F840A293F69EF7F3F840A293F69EF7F3FBF1B7D3F85ED7F3FBF1B7D3F85ED7F3F22F0D03E85EB7F3F22F0D03E85EB7F3FFC370CBF6CE97F3FFC370CBF6CE97F3F6FFD7FBF39E77F3F6FFD7FBF39E77F3FBF6708BFEBE47F3FBF6708BFEBE47F3FFE2DD93E83E27F3FFE2DD93E83E27F3F78BF7D3F01E07F3F78BF7D3F01E07F3F819C253F64DD7F3F819C253F64DD7F3F389395BEAEDA7F3F389395BEAEDA7F3F576D76BFDDD77F3F576D76BFDDD77F3FB3803FBFF2D47F3FB3803FBFF2D47F3F18F41D3EEDD17F3F18F41D3EEDD17F3F8E2C6A3FCECE7F3F8E2C6A3FCECE7F3FAA8F553F94CB7F3FAA8F553F94CB7F3FB78659BC40C87F3FB78659BC40C87F3FE73B59BFD2C47F3FE73B59BFD2C47F3F5F5867BF4AC17F3F5F5867BF4AC17F3FEA0803BEA8BD7F3FEA0803BEA8BD7F3F2DF2433FEBB97F3F2DF2433FEBB97F3FB57F743F15B67F3FB57F743F15B67F3F6C85883E24B27F3F6C85883E24B27F3F74BC2ABF19AE7F3F74BC2ABF19AE7F3F44C27CBFF4A97F3F44C27CBFF4A97F3FE0CACCBEB4A57F3FE0CACCBEB4A57F3FE81B0E3F5BA17F3FE81B0E3F5BA17F3FBBF57F3FE79C7F3FBBF57F3FE79C7F3F807B063F59987F3F807B063F59987F3F6D46DDBEB1937F3F6D46DDBEB1937F3FB3097EBFEF8E7F3FB3097EBFEF8E7F3F80E023BF128A7F3F80E023BF128A7F3F44E7993E1B857F3F44E7993E1B857F3F0308773F0B807F3F0308773F0B807F3FD1FD3D3FE07A7F3FD1FD3D3FE07A7F3F0EE726BE9B757F3F0EE726BE9B757F3F95146BBF3B707F3F95146BBF3B707F3FA64D54BFC26A7F3FA64D54BFC26A7F3F2C43B53C2E657F3F2C43B53C2E657F3FA26C5A3F815F7F3FA26C5A3F815F7F3FAC5D663FB9597F3FAC5D663FB9597F3FB714F43DD7537F3FB714F43DD7537F3F836545BFDB4D7F3F836545BFDB4D7F3F56D173BFC4477F3F56D173BFC4477F3FE62584BE94417F3FE62584BE94417F3FF76A2C3F493B7F3FF76A2C3F493B7F3FB8637C3FE4347F3FB8637C3FE4347F3F83A1C83E662E7F3F83A1C83E662E7F3FFAFC0FBFCD277F3FFAFC0FBFCD277F3FE5E87FBF19217F3FE5E87FBF19217F3F908C04BF4C1A7F3F908C04BF4C1A7F3F6D5AE13E65137F3F6D5AE13E65137F3FD54E7E3F630C7F3FD54E7E3F630C7F3F3521223F48057F3F3521223F48057F3F3A389EBE12FE7E3F3A389EBE12FE7E3FBC9D77BFC2F67E3FBC9D77BFC2F67E3F20773CBF58EF7E3F20773CBF58EF7E3FACD62F3ED4E77E3FACD62F3ED4E77E3FE4F76B3F36E07E3FE4F76B3F36E07E3F6107533F7ED87E3F6107533F7ED87E3F5ABFFDBCABD07E3F5ABFFDBCABD07E3FFC985BBFBFC87E3FFC985BBFBFC87E3F5A5E65BFB8C07E3F5A5E65BFB8C07E3FB512E2BD98B87E3FB512E2BD98B87E3FE4D4463F5DB07E3FE4D4463F5DB07E3F141E733F08A87E3F141E733F08A87E3F72877F3E999F7E3F72877F3E999F7E3F05162EBF11977E3F05162EBF11977E3F1C007CBF6E8E7E3F1C007CBF6E8E7E3F1F74C4BEB0857E3F1F74C4BEB0857E3F28DB113FD97C7E3F28DB113FD97C7E3FEDD67F3FE8737E3FEDD67F3FE8737E3FF69A023FDD6A7E3FF69A023FDD6A7E3FE869E5BEB8617E3FE869E5BEB8617E3FDE8E7EBF78587E3FDE8E7EBF78587E3FAA5E20BF1F4F7E3FAA5E20BF1F4F7E3F0486A23EAB457E3F0486A23EAB457E3F7D2E783F1E3C7E3F7D2E783F1E3C7E3FA8EC3A3F76327E3FA8EC3A3F76327E3FC2C238BEB5287E3FC2C238BEB5287E3F78D66CBFD91E7E3F78D66CBFD91E7E3FE0BC51BFE4147E3FE0BC51BFE4147E3F381B233DD40A7E3F381B233DD40A7E3FEEC05C3FAB007E3FEEC05C3FAB007E3F6E5A643F67F67D3F6E5A643F67F67D3F2A0CD03D09EC7D3F2A0CD03D09EC7D3F494048BF92E17D3F494048BF92E17D3FF16572BF00D77D3FF16572BF00D77D3FF8BD76BE54CC7D3FF8BD76BE54CC7D3F95BD2F3F8FC17D3F95BD2F3F8FC17D3F71977B3FAFB67D3F71977B3FAFB67D3FCB42C03EB6AB7D3FCB42C03EB6AB7D3F6AB613BFA2A07D3F6AB613BFA2A07D3FD4BF7FBF75957D3FD4BF7FBF75957D3FBFA600BF2D8A7D3FBFA600BF2D8A7D3FC874E93ECC7E7D3FC874E93ECC7E7D3FCBC97E3F50737D3FCBC97E3F50737D3FE7981E3FBB677D3FE7981E3FBB677D3F8BD0A6BE0C5C7D3F8BD0A6BE0C5C7D3F44BA78BF42507D3F44BA78BF42507D3F6F5E39BF5F447D3F6F5E39BF5F447D3F25AB413E62387D3F25AB413E62387D3F4CB06D3F4B2C7D3F4CB06D3F4B2C7D3F2A6E503F1A207D3F2A6E503F1A207D3F7E5347BDCF137D3F7E5347BDCF137D3F73E45DBF6A077D3F73E45DBF6A077D3FEE5163BFEBFA7C3FEE5163BFEBFA7C3F7301BEBD53EE7C3F7301BEBD53EE7C3FA9A7493FA0E17C3FA9A7493FA0E17C3FF2A8713FD4D47C3FF2A8713FD4D47C3F8BEF6D3EEEC77C3F8BEF6D3EEEC77C3F9F6131BFEDBA7C3F9F6131BFEDBA7C3FBB297BBFD3AD7C3FBB297BBFD3AD7C3F9C0DBCBE9FA07C3F9C0DBCBE9FA07C3FB68E153F52937C3FB68E153F52937C3F99A37F3FEA857C3F99A37F3FEA857C3FE45FFD3E68787C3FE45FFD3E68787C3FFA7AEDBECD6A7C3FFA7AEDBECD6A7C3F9DFF7EBF185D7C3F9DFF7EBF185D7C3FF6CF1CBF494F7C3FF6CF1CBF494F7C3FBA17AB3E60417C3FBA17AB3E60417C3F0E41793F5D337C3F0E41793F5D337C3F7FCC373F40257C3F7FCC373F40257C3FA48F4ABE0A177C3FA48F4ABE0A177C3F5B856EBFBA087C3F5B856EBFBA087C3F461B4FBF50FA7B3F461B4FBF50FA7B3FC5876B3DCCEB7B3FC5876B3DCCEB7B3F84035F3F2EDD7B3F84035F3F2EDD7B3FDF44623F77CE7B3FDF44623F77CE7B3FEDF2AB3DA6BF7B3FEDF2AB3DA6BF7B3FFD0A4BBFBBB07B3FFD0A4BBFBBB07B3F1AE770BFB6A17B3F1AE770BFB6A17B3F591C65BE97927B3F591C65BE97927B3F1A02333F5F837B3F1A02333F5F837B3FFCB67A3F0D747B3FFCB67A3F0D747B3FA7D4B73EA1647B3FA7D4B73EA1647B3F016417BF1C557B3F016417BF1C557B3F3E827FBF7C457B3F3E827FBF7C457B3F366DF9BEC3357B3F366DF9BEC3357B3F697CF13EF1257B3F697CF13EF1257B3F51307F3F04167B3F51307F3F04167B3FE0031B3FFE057B3FE0031B3FFE057B3F7A5BAFBEDEF57A3F7A5BAFBEDEF57A3FD8C279BFA5E57A3FD8C279BFA5E57A3FDF3636BF51D57A3FDF3636BF51D57A3F1470533EE5C47A3F1470533EE5C47A3FA2556F3F5EB47A3FA2556F3F5EB47A3F3BC44D3FBEA37A3F3BC44D3FBEA37A3FA9DB87BD04937A3FA9DB87BD04937A3F1C1E60BF30827A3F1C1E60BF30827A3F463361BF43717A3F463361BF43717A3FF3E099BD3C607A3FF3E099BD3C607A3F3F6A4C3F1B4F7A3F3F6A4C3F1B4F7A3F6D20703FE13D7A3F6D20703FE13D7A3F8E445C3E8D2C7A3F8E445C3E8D2C7A3FFE9E34BF201B7A3FFE9E34BF201B7A3F353F7ABF99097A3F353F7ABF99097A3F0298B3BEF8F7793F0298B3BEF8F7793F4336193F3EE6793F4336193F3EE6793FC35B7F3F6AD4793FC35B7F3F6AD4793F8875F53E7DC2793F8875F53E7DC2793FFF78F5BE76B0793FFF78F5BE76B0793FE75B7FBF559E793FE75B7FBF559E793FAE3419BF1B8C793FAE3419BF1B8C793FB69BB33EC879793FB69BB33EC879793F9F3F7A3F5A67793F9F3F7A3F5A67793F979D343FD454793F979D343FD454793F464C5CBE3342793F464C5CBE3342793F1C2170BF7A2F793F1C2170BF7A2F793F0F694CBFA61C793F0F694CBFA61C793FB6F0993DB909793FB6F0993DB909793F3634613FB3F6783F3634613FB3F6783F281D603F93E3783F281D603F93E3783FE4CB873D5AD0783FE4CB873D5AD0783F67C54DBF07BD783F67C54DBF07BD783FEF546FBF9BA9783FEF546FBF9BA9783F586853BE1596783F586853BE1596783F4238363F7682783F4238363F7682783F69C2793FBE6E783F69C2793FBE6E783FC357AF3EEB5A783FC357AF3EEB5A783F72051BBF0047783F72051BBF0047783F29307FBFFB32783F29307FBFFB32783FED78F1BEDD1E783FED78F1BEDD1E783FAA70F93EA50A783FAA70F93EA50A783F5D827F3F54F6773F5D827F3F54F6773F6962173FE9E1773F6962173FE9E1773F57D8B7BE65CD773F57D8B7BE65CD773F62B77ABFC8B8773F62B77ABFC8B8773FB00033BF11A4773FB00033BF11A4773F0D24653E418F773F0D24653E418F773FC5E7703F587A773FC5E7703F587A773FC9094B3F5565773FC9094B3F5565773FAC02ACBD3950773FAC02ACBD3950773FCB4562BF043B773FCB4562BF043B773F8C025FBFB525773F8C025FBFB525773F36686BBD4D10773F36686BBD4D10773F6F1C4F3FCCFA763F6F1C4F3FCCFA763FA4846E3F31E5763FA4846E3F31E5763FE5874A3E7ECF763FE5874A3E7ECF763FDFCD37BFB1B9763FDFCD37BFB1B9763F9A4079BFCAA3763F9A4079BFCAA3763F0014ABBECA8D763F0014ABBECA8D763F86D11C3FB277763F86D11C3FB277763F70FF7E3F7F61763F70FF7E3F7F61763F7A77ED3E344B763F7A77ED3E344B763F5363FDBECF34763F5363FDBECF34763FB4A37FBF521E763FB4A37FBF521E763F1B8D15BFBB07763F1B8D15BFBB07763F4911BC3E0AF1753F4911BC3E0AF1753F1D2A7B3F41DA753F1D2A7B3F41DA753F3260313F5FC3753F3260313F5FC3753F3BF76DBE63AC753F3BF76DBE63AC753F99A971BF4E95753F99A971BF4E95753F71A649BF207E753F71A649BF207E753F3011BE3DD966753F3011BE3DD966753FD752633F784F753FD752633F784F753F76E35D3FFF37753F76E35D3FFF37753FEC33473D6D20753FEC33473D6D20753F506F50BFC108753F506F50BFC108753F90AF6DBFFCF0743F90AF6DBFFCF0743F62A341BE1ED9743F62A341BE1ED9743FCC5F393F28C1743FCC5F393F28C1743FCCB9783F18A9743FCCB9783F18A9743FCFCCA63EEF90743FCFCCA63EEF90743F749A1EBFAD78743F749A1EBFAD78743F9AC97EBF5260743F9AC97EBF5260743F4471E9BEDD47743F4471E9BEDD47743F74A8003F502F743F74A8003F502F743FEABF7F3FAA16743FEABF7F3FAA16743FCDB4133FEBFD733FCDB4133FEBFD733F7546C0BE13E5733F7546C0BE13E5733FCF977BBF22CC733FCF977BBF22CC733F25BC2FBF18B3733F25BC2FBF18B3733FA3C5763EF599733FA3C5763EF599733F9466723FB980733F9466723FB980733F0E3F483F6467733F0E3F483F6467733FE31BD0BDF64D733FE31BD0BDF64D733F535B64BF7034733F535B64BF7034733FEEBF5CBFD01A733FEEBF5CBFD01A733F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F07D7A33AA46A573F07D7A33AB7C7683FFFD6233BB7C7683FFFD6233BC381103E6AC2753BC381103E6AC2753BCFBD41BFDED6A33BCFBD41BFDED6A33B107C75BF75CCCC3B107C75BF75CCCC3B8C0F8FBEF8C1F53B8C0F8FBEF8C1F53B4630283FB15B0F3C4630283FB15B0F3C95467D3F57D6233C95467D3F57D6233C3201D33EED50383C3201D33EED50383CF8440BBF6FCB4C3CF8440BBF6FCB4C3C5CFF7FBFDD45613C5CFF7FBFDD45613CD85C09BF33C0753CD85C09BF33C0753C2220D73E381D853C2220D73E381D853C72987D3F495A8F3C72987D3F495A8F3C4479263F4C97993C4479263F4C97993C106893BE3ED4A33C106893BE3ED4A33C251E76BF2011AE3C251E76BF2011AE3CB34040BFF04DB83CB34040BFF04DB83C6979193EAD8AC23C6979193EAD8AC23CC8B6693F57C7CC3CC8B6693F57C7CC3C102F563FEB03D73C102F563FEB03D73C150511BC6940E13C150511BC6940E13CE7A158BFD07CEB3CE7A158BFD07CEB3CFCD367BF20B9F53CFCD367BF20B9F53C388707BE55F5FF3C388707BE55F5FF3C0837433FB918053D0837433FB918053D0ED5743FB9360A3D0ED5743FB9360A3D2AB48A3EAB540F3D2AB48A3EAB540F3DE9E329BF8E72143DE9E329BF8E72143DA4EF7CBF6290193DA4EF7CBF6290193D06DECEBE27AE1E3D06DECEBE27AE1E3D4D2A0D3FDBCB233D4D2A0D3FDBCB233D39FA7F3F7FE9283D39FA7F3F7FE9283DF771073F11072E3DF771073F11072E3DC23ADBBE9124333DC23ADBBE9124333D38E57DBFFF41383D38E57DBFFF41383DEABE24BF5A5F3D3DEABE24BF5A5F3D3D9FBD973EA27C423D9FBD973EA27C423D4BBB763FD699473D4BBB763FD699473DBCBF3E3FF5B64C3DBCBF3E3FF5B64C3DFB6D22BE00D4513DFB6D22BE00D4513D28A16ABFF4F0563D28A16ABFF4F0563D31EF54BFD30D5C3D31EF54BFD30D5C3DA103913C9C2A613DA103913C9C2A613DD0D4593F4D47663DD0D4593F4D47663D9ADB663FE7636B3D9ADB663FE7636B3DE813FD3D6880703DE813FD3D6880703D56AC44BFD19C753D56AC44BFD19C753D222974BF21B97A3D222974BF21B97A3DFF5586BE57D57F3DFF5586BE57D57F3D24942B3FBA78823D24942B3FBA78823DA0937C3FBA06853DA0937C3FBA06853DB4B6CA3EAD94873DB4B6CA3EAD94873DCD0C0FBF92228A3DCD0C0FBF92228A3DF4EF7FBF69B08C3DF4EF7FBF69B08C3D5E8405BF323E8F3D5E8405BF323E8F3DFD50DF3EECCB913DFD50DF3EECCB913DE72C7E3F9759943DE72C7E3F9759943D4301233F33E7963D4301233F33E7963D23109CBEBF74993D23109CBEBF74993D7F5377BF3C029C3D7F5377BF3C029C3DF23A3DBFA88F9E3DF23A3DBFA88F9E3D4B5F2B3E051DA13D4B5F2B3E051DA13DD4866B3F51AAA33DD4866B3F51AAA33D0BAB533F8C37A63D0BAB533F8C37A63DCE81D9BCB6C4A83DCE81D9BCB6C4A83D5B035BBFCF51AB3D5B035BBFCF51AB3D97DE65BFD6DEAD3D97DE65BFD6DEAD3D4D14EBBDCC6BB03D4D14EBBDCC6BB03DB31D463FAFF8B23DB31D463FAFF8B23D5178733F8085B53D5178733F8085B53D23F5813E3F12B83D23F5813E3F12B83DED402DBFEA9EBA3DED402DBFEA9EBA3D8C327CBF832BBD3D8C327CBF832BBD3D518BC6BE08B8BF3D518BC6BE08B8BF3D6EEC103F7A44C23D6EEC103F7A44C23D8EE07F3FD7D0C43D8EE07F3FD7D0C43D1894033F215DC73D1894033F215DC73DBC62E3BE56E9C93DBC62E3BE56E9C93D7D6F7EBF7675CC3D7D6F7EBF7675CC3D574021BF8201CF3D574021BF8201CF3D865FA03E788DD13D865FA03E788DD13DBCE6773F5919D43DBCE6773F5919D43D5CB23B3F24A5D63D5CB23B3F24A5D63D2B4D34BED930D93D2B4D34BED930D93DC6676CBF78BCDB3DC6676CBF78BCDB3DA76252BF0048DE3DA76252BF0048DE3DD0FD103D72D3E03DD0FD103D72D3E03D822D5C3FCD5EE33D822D5C3FCD5EE33DF7DC643F10EAE53DF7DC643F10EAE53DFB0FD93D3C75E83DFB0FD93D3C75E83D178B47BF5000EB3D178B47BF5000EB3D9EC272BF4C8BED3D9EC272BF4C8BED3D56237BBE2F16F03D56237BBE2F16F03D3DEA2E3FFAA0F23D3DEA2E3FFAA0F23D68CC7B3FAC2BF53D68CC7B3FAC2BF53DF25BC23E46B6F73DF25BC23E46B6F73D28C912BFC540FA3D28C912BFC540FA3D05CC7FBF2BCBFC3D05CC7FBF2BCBFC3D2EA101BF7755FF3D2EA101BF7755FF3DED6FE73ED5EF003EED6FE73ED5EF003EF8AC7E3FE134023EF8AC7E3FE134023E2F7C1F3FDF79033E2F7C1F3FDF79033EB1ABA4BED0BE043EB1ABA4BED0BE043E007578BFB303063E007578BFB303063E03263ABF8948073E03263ABF8948073E6D373D3E518D083E6D373D3E518D083EFA436D3F0BD2093EFA436D3F0BD2093E0B16513FB7160B3E0B16513FB7160B3ECF3735BD545B0C3ECF3735BD545B0C3E3E535DBFE39F0D3E3E535DBFE39F0D3EC0D663BF64E40E3EC0D663BF64E40E3E4F07C7BDD628103E4F07C7BDD628103E7AF4483F396D113E7AF4483F396D113E0D08723F8DB1123E0D08723F8DB1123E5D57723ED3F5133E5D57723ED3F5133E0B9030BF093A153E0B9030BF093A153E38617BBF307E163E38617BBF307E163EAE28BEBE47C2173EAE28BEBE47C2173EEFA2143F4F06193EEFA2143F4F06193E5BB27F3F484A1A3E5BB27F3F484A1A3E5557FF3E308E1B3E5557FF3E308E1B3E7878EBBE09D21C3E7878EBBE09D21C3E58E57EBFD1151E3E58E57EBFD1151E3ED4B41DBF89591F3ED4B41DBF89591F3E8FF4A83E319D203E8FF4A83E319D203E49FE783FC9E0213E49FE783FC9E0213EED95383F5024233EED95383F5024233EE41D46BEC667243EE41D46BEC667243E6C1B6EBF2CAB253E6C1B6EBF2CAB253E3DC54FBF80EE263E3DC54FBF80EE263E2D6E593DC331283E2D6E593DC331283E8A745E3FF574293E8A745E3FF574293EF8CB623F16B82A3EF8CB623F16B82A3EA4FAB43D26FB2B3EA4FAB43D26FB2B3ED5594ABF233E2D3ED5594ABF233E2D3EA14871BF0F812E3EA14871BF0F812E3E888669BEE9C32F3E888669BEE9C32F3E4F32323FB106313E4F32323FB106313EFCF07A3F6749323EFCF07A3F6749323E99F1B93E0B8C333E99F1B93E0B8C333EBC7916BF9CCE343EBC7916BF9CCE343E90937FBF1B11363E90937FBF1B11363E2F67FBBE8753373E2F67FBBE8753373E4B7CEF3EE195383E4B7CEF3EE195383E9B187F3F27D8393E9B187F3F27D8393E4FEA1B3F5B1A3B3E4FEA1B3F5B1A3B3E093AADBE7B5C3C3E093AADBE7B5C3C3E938279BF889E3D3E938279BF889E3D3E240237BF81E03E3E240237BF81E03E3E61004F3E6822403E61004F3E6822403E18EE6E3F3A64413E18EE6E3F3A64413E45704E3FF8A5423E45704E3FF8A5423E2EA07DBDA3E7433E2EA07DBDA3E7433E60915FBF3A29453E60915FBF3A29453EA3BC61BFBC6A463EA3BC61BFBC6A463E59EAA2BD2AAC473E59EAA2BD2AAC473E21BB4B3F84ED483E21BB4B3F84ED483E5D84703FC92E4A3E5D84703FC92E4A3E04B1603EF96F4B3E04B1603EF96F4B3EFFD033BF15B14C3EFFD033BF15B14C3EB97B7ABF1BF24D3EB97B7ABF1BF24D3EC9B6B5BE0D334F3EC9B6B5BE0D334F3E844D183FE973503E844D183FE973503EA46F7F3FB0B4513EA46F7F3FB0B4513EFE71F73E61F5523EFE71F73E61F5523E507BF3BEFD35543E507BF3BEFD35543EC0467FBF8376553EC0467FBF8376553EAA1C1ABFF3B6563EAA1C1ABFF3B6563E0A7CB13E4EF7573E0A7CB13E4EF7573EDC017A3F9237593EDC017A3F9237593EAF6A353FC0775A3EAF6A353FC0775A3EB7DE57BED7B75B3EB7DE57BED7B75B3EF9BB6FBFD8F75C3EF9BB6FBFD8F75C3E28174DBFC2375E3E28174DBFC2375E3E8CE6903D96775F3E8CE6903D96775F3EB9A9603F53B7603EB9A9603F53B7603EC7A8603FF9F6613EC7A8603FF9F6613EC9D6903D8736633EC9D6903D8736633E57184DBFFF75643E57184DBFFF75643E47BB6FBF5EB5653E47BB6FBF5EB5653EFED657BEA7F4663EFED657BEA7F4663E146C353FD833683E146C353FD833683E6F017A3FF172693E6F017A3FF172693E5578B13EF2B16A3E5578B13EF2B16A3E3E1E1ABFDBF06B3E3E1E1ABFDBF06B3E9A467FBFAC2F6D3E9A467FBFAC2F6D3ED777F3BE646E6E3ED777F3BE646E6E3E7375F73E04AD6F3E7375F73E04AD6F3EC66F7F3F8CEB703EC66F7F3F8CEB703EED4B183FFB29723EED4B183FFB29723E7BBAB5BE5168733E7BBAB5BE5168733E217C7ABF8EA6743E217C7ABF8EA6743E97CF33BFB3E4753E97CF33BFB3E4753EB9B8603EBE22773EB9B8603EBE22773E0B85703FAF60783E0B85703FAF60783EEFB94B3F889E793EEFB94B3F889E793E1AFAA2BD46DC7A3E1AFAA2BD46DC7A3E92BD61BFEB197C3E92BD61BFEB197C3E69905FBF77577D3E69905FBF77577D3EA2807DBDE8947E3EA2807DBDE8947E3E70714E3F3FD27F3E70714E3F3FD27F3E62ED6E3FBE87803E62ED6E3FBE87803EA4F84E3E4F26813EA4F84E3E4F26813E860337BFD4C4813E860337BFD4C4813E228279BF4A63823E228279BF4A63823E5136ADBEB401833E5136ADBEB401833EE0EB1B3F10A0833EE0EB1B3F10A0833E70187F3F5E3E843E70187F3F5E3E843ECD78EF3EA0DC843ECD78EF3EA0DC843EA06AFBBED37A853EA06AFBBED37A853EAD937FBFF918863EAD937FBFF918863E237816BF11B7863E237816BF11B7863E47F5B93E1B55873E47F5B93E1B55873E61F17A3F17F3873E61F17A3F17F3873EE430323F0691883EE430323F0691883E3A8E69BEE62E893E3A8E69BEE62E893E4A4971BFB9CC893E4A4971BFB9CC893E9F584ABF7D6A8A3E9F584ABF7D6A8A3E620AB53D33088B3E620AB53D33088B3EE3CC623FDBA58B3EE3CC623FDBA58B3E90735E3F75438C3E90735E3F75438C3E9D4E593D00E18C3E9D4E593D00E18C3E65C64FBF7D7E8D3E65C64FBF7D7E8D3EB31A6EBFEB1B8E3EB31A6EBFEB1B8E3E231646BE4BB98E3E231646BE4BB98E3E4C97383F9C568F3E4C97383F9C568F3ED3FD783FDFF38F3ED3FD783FDFF38F3ED4F0A83E1291903ED4F0A83E1291903E62B61DBF372E913E62B61DBF372E913E29E57EBF4DCB913E29E57EBF4DCB913EF674EBBE5468923EF674EBBE5468923EC15AFF3E4D05933EC15AFF3E4D05933E73B27F3F36A2933E73B27F3F36A2933E53A1143F103F943E53A1143F103F943E592CBEBEDBDB943E592CBEBEDBDB943E97617BBF9678953E97617BBF9678953E9D8E30BF4215963E9D8E30BF4215963E0B5F723EDFB1963E0B5F723EDFB1963EB208723F6D4E973EB208723F6D4E973E40F3483FEBEA973E40F3483FEBEA973E0917C7BD5987983E0917C7BD5987983EA7D763BFB823993EA7D763BFB823993E40525DBF07C0993E40525DBF07C0993E3C1835BD465C9A3E3C1835BD465C9A3E2F17513F76F89A3E2F17513F76F89A3E3C436D3F96949B3E3C436D3F96949B3EA92F3D3EA5309C3EA92F3D3EA5309C3E5E273ABFA5CC9C3E5E273ABFA5CC9C3E867478BF95689D3E867478BF95689D3EF3A7A4BE74049E3EF3A7A4BE74049E3EBA7D1F3F44A09E3EBA7D1F3F44A09E3EC5AC7E3F033C9F3EC5AC7E3F033C9F3E666CE73EB2D79F3E666CE73EB2D79F3EE2A201BF5073A03EE2A201BF5073A03E"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x3727C5AC> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.dynamic_slice %arg0, %arg27, %9, sizes = [1, 8] : (tensor<10x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %13 = stablehlo.reshape %12 : (tensor<1x8xf32>) -> tensor<8xf32> + %14 = stablehlo.dynamic_slice %0, %arg28, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %15 = stablehlo.reshape %14 : (tensor<1x4xf32>) -> tensor<4xf32> + %16 = stablehlo.dynamic_slice %1, %arg28, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %17 = stablehlo.reshape %16 : (tensor<1x4xf32>) -> tensor<4xf32> + %18 = stablehlo.iota dim = 0 : tensor<256xi32> + %19 = stablehlo.broadcast_in_dim %arg29, dims = [] : (tensor) -> tensor<256xi32> + %20 = stablehlo.compare LE, %18, %19, SIGNED : (tensor<256xi32>, tensor<256xi32>) -> tensor<256xi1> + %21 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256xf32> + %22 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256xf32> + %23 = stablehlo.select %20, %21, %22 : tensor<256xi1>, tensor<256xf32> + %24 = stablehlo.multiply %13, %13 : tensor<8xf32> + %25 = stablehlo.reduce(%24 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %26 = stablehlo.divide %25, %7 : tensor + %27 = stablehlo.add %26, %6 : tensor + %28 = stablehlo.rsqrt %27 : tensor + %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor) -> tensor<8xf32> + %30 = stablehlo.multiply %13, %29 : tensor<8xf32> + %31 = stablehlo.multiply %30, %arg5 : tensor<8xf32> + %32 = stablehlo.dot_general %31, %arg10, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %33 = stablehlo.add %32, %arg13 : tensor<8xf32> + %34 = stablehlo.reshape %33 : (tensor<8xf32>) -> tensor<2x4xf32> + %35 = stablehlo.dot_general %31, %arg8, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %36 = stablehlo.add %35, %arg12 : tensor<4xf32> + %37 = stablehlo.reshape %36 : (tensor<4xf32>) -> tensor<1x4xf32> + %38 = stablehlo.dot_general %31, %arg11, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %39 = stablehlo.add %38, %arg14 : tensor<4xf32> + %40 = stablehlo.reshape %39 : (tensor<4xf32>) -> tensor<1x4xf32> + %41 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %42 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %43 = stablehlo.multiply %34, %41 : tensor<2x4xf32> + %44 = stablehlo.slice %34 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %45 = stablehlo.slice %34 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %46 = stablehlo.negate %45 : tensor<2x2xf32> + %47 = stablehlo.concatenate %46, %44, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %48 = stablehlo.multiply %47, %42 : tensor<2x4xf32> + %49 = stablehlo.add %43, %48 : tensor<2x4xf32> + %50 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %51 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %52 = stablehlo.multiply %37, %50 : tensor<1x4xf32> + %53 = stablehlo.slice %37 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %54 = stablehlo.slice %37 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %55 = stablehlo.negate %54 : tensor<1x2xf32> + %56 = stablehlo.concatenate %55, %53, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %57 = stablehlo.multiply %56, %51 : tensor<1x4xf32> + %58 = stablehlo.add %52, %57 : tensor<1x4xf32> + %59 = stablehlo.reshape %58 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %60 = stablehlo.dynamic_update_slice %arg30, %59, %10, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %61 = stablehlo.reshape %40 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %62 = stablehlo.dynamic_update_slice %arg31, %61, %10, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %63 = stablehlo.slice %60 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %64 = stablehlo.reshape %63 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %65 = stablehlo.slice %62 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %66 = stablehlo.reshape %65 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %67 = stablehlo.reshape %49 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %68 = stablehlo.dot_general %67, %64, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %69 = stablehlo.reshape %68 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %70 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %71 = stablehlo.multiply %69, %70 : tensor<2x256xf32> + %72 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %73 = stablehlo.add %71, %72 : tensor<2x256xf32> + %74 = stablehlo.reduce(%73 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %75 = stablehlo.broadcast_in_dim %74, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %76 = stablehlo.subtract %73, %75 : tensor<2x256xf32> + %77 = stablehlo.exponential %76 : tensor<2x256xf32> + %78 = stablehlo.reduce(%77 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %79 = stablehlo.broadcast_in_dim %78, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %80 = stablehlo.divide %77, %79 : tensor<2x256xf32> + %81 = stablehlo.reshape %80 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %82 = stablehlo.dot_general %81, %66, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %83 = stablehlo.reshape %82 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %84 = stablehlo.reshape %83 : (tensor<2x4xf32>) -> tensor<8xf32> + %85 = stablehlo.dot_general %84, %arg9, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %86 = stablehlo.add %13, %85 : tensor<8xf32> + %87 = stablehlo.multiply %86, %86 : tensor<8xf32> + %88 = stablehlo.reduce(%87 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %89 = stablehlo.divide %88, %7 : tensor + %90 = stablehlo.add %89, %6 : tensor + %91 = stablehlo.rsqrt %90 : tensor + %92 = stablehlo.broadcast_in_dim %91, dims = [] : (tensor) -> tensor<8xf32> + %93 = stablehlo.multiply %86, %92 : tensor<8xf32> + %94 = stablehlo.multiply %93, %arg6 : tensor<8xf32> + %95 = stablehlo.dot_general %94, %arg4, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %96 = stablehlo.dot_general %94, %arg7, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %97 = stablehlo.negate %95 : tensor<16xf32> + %98 = stablehlo.exponential %97 : tensor<16xf32> + %99 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %100 = stablehlo.add %99, %98 : tensor<16xf32> + %101 = stablehlo.divide %99, %100 : tensor<16xf32> + %102 = stablehlo.multiply %95, %101 : tensor<16xf32> + %103 = stablehlo.multiply %102, %96 : tensor<16xf32> + %104 = stablehlo.dot_general %103, %arg3, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %105 = stablehlo.add %86, %104 : tensor<8xf32> + %106 = stablehlo.multiply %105, %105 : tensor<8xf32> + %107 = stablehlo.reduce(%106 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %108 = stablehlo.divide %107, %7 : tensor + %109 = stablehlo.add %108, %6 : tensor + %110 = stablehlo.rsqrt %109 : tensor + %111 = stablehlo.broadcast_in_dim %110, dims = [] : (tensor) -> tensor<8xf32> + %112 = stablehlo.multiply %105, %111 : tensor<8xf32> + %113 = stablehlo.multiply %112, %arg17 : tensor<8xf32> + %114 = stablehlo.dot_general %113, %arg22, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %115 = stablehlo.add %114, %arg25 : tensor<8xf32> + %116 = stablehlo.reshape %115 : (tensor<8xf32>) -> tensor<2x4xf32> + %117 = stablehlo.dot_general %113, %arg20, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %118 = stablehlo.add %117, %arg24 : tensor<4xf32> + %119 = stablehlo.reshape %118 : (tensor<4xf32>) -> tensor<1x4xf32> + %120 = stablehlo.dot_general %113, %arg23, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %121 = stablehlo.add %120, %arg26 : tensor<4xf32> + %122 = stablehlo.reshape %121 : (tensor<4xf32>) -> tensor<1x4xf32> + %123 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %124 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %125 = stablehlo.multiply %116, %123 : tensor<2x4xf32> + %126 = stablehlo.slice %116 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %127 = stablehlo.slice %116 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %128 = stablehlo.negate %127 : tensor<2x2xf32> + %129 = stablehlo.concatenate %128, %126, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %130 = stablehlo.multiply %129, %124 : tensor<2x4xf32> + %131 = stablehlo.add %125, %130 : tensor<2x4xf32> + %132 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %133 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %134 = stablehlo.multiply %119, %132 : tensor<1x4xf32> + %135 = stablehlo.slice %119 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %136 = stablehlo.slice %119 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %137 = stablehlo.negate %136 : tensor<1x2xf32> + %138 = stablehlo.concatenate %137, %135, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %139 = stablehlo.multiply %138, %133 : tensor<1x4xf32> + %140 = stablehlo.add %134, %139 : tensor<1x4xf32> + %141 = stablehlo.reshape %140 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %142 = stablehlo.dynamic_update_slice %60, %141, %11, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %143 = stablehlo.reshape %122 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %144 = stablehlo.dynamic_update_slice %62, %143, %11, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %145 = stablehlo.slice %142 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %146 = stablehlo.reshape %145 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %147 = stablehlo.slice %144 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %148 = stablehlo.reshape %147 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %149 = stablehlo.reshape %131 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %150 = stablehlo.dot_general %149, %146, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %151 = stablehlo.reshape %150 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %152 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %153 = stablehlo.multiply %151, %152 : tensor<2x256xf32> + %154 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %155 = stablehlo.add %153, %154 : tensor<2x256xf32> + %156 = stablehlo.reduce(%155 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %157 = stablehlo.broadcast_in_dim %156, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %158 = stablehlo.subtract %155, %157 : tensor<2x256xf32> + %159 = stablehlo.exponential %158 : tensor<2x256xf32> + %160 = stablehlo.reduce(%159 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %161 = stablehlo.broadcast_in_dim %160, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %162 = stablehlo.divide %159, %161 : tensor<2x256xf32> + %163 = stablehlo.reshape %162 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %164 = stablehlo.dot_general %163, %148, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %165 = stablehlo.reshape %164 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %166 = stablehlo.reshape %165 : (tensor<2x4xf32>) -> tensor<8xf32> + %167 = stablehlo.dot_general %166, %arg21, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %168 = stablehlo.add %105, %167 : tensor<8xf32> + %169 = stablehlo.multiply %168, %168 : tensor<8xf32> + %170 = stablehlo.reduce(%169 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %171 = stablehlo.divide %170, %7 : tensor + %172 = stablehlo.add %171, %6 : tensor + %173 = stablehlo.rsqrt %172 : tensor + %174 = stablehlo.broadcast_in_dim %173, dims = [] : (tensor) -> tensor<8xf32> + %175 = stablehlo.multiply %168, %174 : tensor<8xf32> + %176 = stablehlo.multiply %175, %arg18 : tensor<8xf32> + %177 = stablehlo.dot_general %176, %arg16, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %178 = stablehlo.dot_general %176, %arg19, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %179 = stablehlo.negate %177 : tensor<16xf32> + %180 = stablehlo.exponential %179 : tensor<16xf32> + %181 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %182 = stablehlo.add %181, %180 : tensor<16xf32> + %183 = stablehlo.divide %181, %182 : tensor<16xf32> + %184 = stablehlo.multiply %177, %183 : tensor<16xf32> + %185 = stablehlo.multiply %184, %178 : tensor<16xf32> + %186 = stablehlo.dot_general %185, %arg15, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %187 = stablehlo.add %168, %186 : tensor<8xf32> + %188 = stablehlo.multiply %187, %187 : tensor<8xf32> + %189 = stablehlo.reduce(%188 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %190 = stablehlo.divide %189, %7 : tensor + %191 = stablehlo.add %190, %6 : tensor + %192 = stablehlo.rsqrt %191 : tensor + %193 = stablehlo.broadcast_in_dim %192, dims = [] : (tensor) -> tensor<8xf32> + %194 = stablehlo.multiply %187, %193 : tensor<8xf32> + %195 = stablehlo.multiply %194, %arg1 : tensor<8xf32> + %196 = stablehlo.dot_general %195, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %197 = stablehlo.constant dense<0> : tensor + %198 = stablehlo.constant dense<0xFF800000> : tensor + %199 = stablehlo.iota dim = 0 : tensor<10xi32> + %200:2 = stablehlo.reduce(%196 init: %198), (%199 init: %197) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %201 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %202 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %203 = stablehlo.or %201, %202 : tensor + %204 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %205 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %206 = stablehlo.and %204, %205 : tensor + %207 = stablehlo.or %203, %206 : tensor + %208 = stablehlo.select %203, %amv_l, %amv_r : tensor, tensor + %209 = stablehlo.select %207, %ami_l, %ami_r : tensor, tensor + stablehlo.return %208, %209 : tensor, tensor + } + return %200#1, %142, %144 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/mimo/prefill.mlir b/src/lib/mlxcel-xla/assets/mimo/prefill.mlir new file mode 100644 index 000000000..046059ba6 --- /dev/null +++ b/src/lib/mlxcel-xla/assets/mimo/prefill.mlir @@ -0,0 +1,234 @@ +module @prefill { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg4: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg5: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg6: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg7: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg10: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg11: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg12: tensor<4xf32> loc("params['layers'][0]['bk']"), %arg13: tensor<8xf32> loc("params['layers'][0]['bq']"), %arg14: tensor<4xf32> loc("params['layers'][0]['bv']"), %arg15: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg16: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg17: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg18: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg19: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg20: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg21: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg22: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg23: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg24: tensor<4xf32> loc("params['layers'][1]['bk']"), %arg25: tensor<8xf32> loc("params['layers'][1]['bq']"), %arg26: tensor<4xf32> loc("params['layers'][1]['bv']"), %arg27: tensor<256xi32> loc("tokens"), %arg28: tensor<256xi32> loc("positions"), %arg29: tensor loc("real_len")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3FF3FF7F3F40510A3FF3FF7F3F3311D5BECCFF7F3F3311D5BECCFF7F3F26707DBF8AFF7F3F26707DBF8AFF7F3F305527BF2EFF7F3F305527BF2EFF7F3F2C3C913EB8FE7F3F2C3C913EB8FE7F3FB8CD753F28FE7F3FB8CD753F28FE7F3FBDFF403F7EFD7F3FBDFF403F7EFD7F3FF6FD14BEB9FC7F3FF6FD14BEB9FC7F3FD53F69BFDAFB7F3FD53F69BFDAFB7F3F64CD56BFE1FA7F3F64CD56BFE1FA7F3F7205913BCEF97F3F7205913BCEF97F3FD006583FA1F87F3FD006583FA1F87F3F6F4E683F59F77F3F6F4E683F59F77F3FD7040C3EF7F57F3FD7040C3EF7F57F3FE87A42BF7BF47F3FE87A42BF7BF47F3F2C2975BFE5F27F3F2C2975BFE5F27F3F36E28CBE34F17F3F36E28CBE34F17F3F840A293F69EF7F3F840A293F69EF7F3FBF1B7D3F85ED7F3FBF1B7D3F85ED7F3F22F0D03E85EB7F3F22F0D03E85EB7F3FFC370CBF6CE97F3FFC370CBF6CE97F3F6FFD7FBF39E77F3F6FFD7FBF39E77F3FBF6708BFEBE47F3FBF6708BFEBE47F3FFE2DD93E83E27F3FFE2DD93E83E27F3F78BF7D3F01E07F3F78BF7D3F01E07F3F819C253F64DD7F3F819C253F64DD7F3F389395BEAEDA7F3F389395BEAEDA7F3F576D76BFDDD77F3F576D76BFDDD77F3FB3803FBFF2D47F3FB3803FBFF2D47F3F18F41D3EEDD17F3F18F41D3EEDD17F3F8E2C6A3FCECE7F3F8E2C6A3FCECE7F3FAA8F553F94CB7F3FAA8F553F94CB7F3FB78659BC40C87F3FB78659BC40C87F3FE73B59BFD2C47F3FE73B59BFD2C47F3F5F5867BF4AC17F3F5F5867BF4AC17F3FEA0803BEA8BD7F3FEA0803BEA8BD7F3F2DF2433FEBB97F3F2DF2433FEBB97F3FB57F743F15B67F3FB57F743F15B67F3F6C85883E24B27F3F6C85883E24B27F3F74BC2ABF19AE7F3F74BC2ABF19AE7F3F44C27CBFF4A97F3F44C27CBFF4A97F3FE0CACCBEB4A57F3FE0CACCBEB4A57F3FE81B0E3F5BA17F3FE81B0E3F5BA17F3FBBF57F3FE79C7F3FBBF57F3FE79C7F3F807B063F59987F3F807B063F59987F3F6D46DDBEB1937F3F6D46DDBEB1937F3FB3097EBFEF8E7F3FB3097EBFEF8E7F3F80E023BF128A7F3F80E023BF128A7F3F44E7993E1B857F3F44E7993E1B857F3F0308773F0B807F3F0308773F0B807F3FD1FD3D3FE07A7F3FD1FD3D3FE07A7F3F0EE726BE9B757F3F0EE726BE9B757F3F95146BBF3B707F3F95146BBF3B707F3FA64D54BFC26A7F3FA64D54BFC26A7F3F2C43B53C2E657F3F2C43B53C2E657F3FA26C5A3F815F7F3FA26C5A3F815F7F3FAC5D663FB9597F3FAC5D663FB9597F3FB714F43DD7537F3FB714F43DD7537F3F836545BFDB4D7F3F836545BFDB4D7F3F56D173BFC4477F3F56D173BFC4477F3FE62584BE94417F3FE62584BE94417F3FF76A2C3F493B7F3FF76A2C3F493B7F3FB8637C3FE4347F3FB8637C3FE4347F3F83A1C83E662E7F3F83A1C83E662E7F3FFAFC0FBFCD277F3FFAFC0FBFCD277F3FE5E87FBF19217F3FE5E87FBF19217F3F908C04BF4C1A7F3F908C04BF4C1A7F3F6D5AE13E65137F3F6D5AE13E65137F3FD54E7E3F630C7F3FD54E7E3F630C7F3F3521223F48057F3F3521223F48057F3F3A389EBE12FE7E3F3A389EBE12FE7E3FBC9D77BFC2F67E3FBC9D77BFC2F67E3F20773CBF58EF7E3F20773CBF58EF7E3FACD62F3ED4E77E3FACD62F3ED4E77E3FE4F76B3F36E07E3FE4F76B3F36E07E3F6107533F7ED87E3F6107533F7ED87E3F5ABFFDBCABD07E3F5ABFFDBCABD07E3FFC985BBFBFC87E3FFC985BBFBFC87E3F5A5E65BFB8C07E3F5A5E65BFB8C07E3FB512E2BD98B87E3FB512E2BD98B87E3FE4D4463F5DB07E3FE4D4463F5DB07E3F141E733F08A87E3F141E733F08A87E3F72877F3E999F7E3F72877F3E999F7E3F05162EBF11977E3F05162EBF11977E3F1C007CBF6E8E7E3F1C007CBF6E8E7E3F1F74C4BEB0857E3F1F74C4BEB0857E3F28DB113FD97C7E3F28DB113FD97C7E3FEDD67F3FE8737E3FEDD67F3FE8737E3FF69A023FDD6A7E3FF69A023FDD6A7E3FE869E5BEB8617E3FE869E5BEB8617E3FDE8E7EBF78587E3FDE8E7EBF78587E3FAA5E20BF1F4F7E3FAA5E20BF1F4F7E3F0486A23EAB457E3F0486A23EAB457E3F7D2E783F1E3C7E3F7D2E783F1E3C7E3FA8EC3A3F76327E3FA8EC3A3F76327E3FC2C238BEB5287E3FC2C238BEB5287E3F78D66CBFD91E7E3F78D66CBFD91E7E3FE0BC51BFE4147E3FE0BC51BFE4147E3F381B233DD40A7E3F381B233DD40A7E3FEEC05C3FAB007E3FEEC05C3FAB007E3F6E5A643F67F67D3F6E5A643F67F67D3F2A0CD03D09EC7D3F2A0CD03D09EC7D3F494048BF92E17D3F494048BF92E17D3FF16572BF00D77D3FF16572BF00D77D3FF8BD76BE54CC7D3FF8BD76BE54CC7D3F95BD2F3F8FC17D3F95BD2F3F8FC17D3F71977B3FAFB67D3F71977B3FAFB67D3FCB42C03EB6AB7D3FCB42C03EB6AB7D3F6AB613BFA2A07D3F6AB613BFA2A07D3FD4BF7FBF75957D3FD4BF7FBF75957D3FBFA600BF2D8A7D3FBFA600BF2D8A7D3FC874E93ECC7E7D3FC874E93ECC7E7D3FCBC97E3F50737D3FCBC97E3F50737D3FE7981E3FBB677D3FE7981E3FBB677D3F8BD0A6BE0C5C7D3F8BD0A6BE0C5C7D3F44BA78BF42507D3F44BA78BF42507D3F6F5E39BF5F447D3F6F5E39BF5F447D3F25AB413E62387D3F25AB413E62387D3F4CB06D3F4B2C7D3F4CB06D3F4B2C7D3F2A6E503F1A207D3F2A6E503F1A207D3F7E5347BDCF137D3F7E5347BDCF137D3F73E45DBF6A077D3F73E45DBF6A077D3FEE5163BFEBFA7C3FEE5163BFEBFA7C3F7301BEBD53EE7C3F7301BEBD53EE7C3FA9A7493FA0E17C3FA9A7493FA0E17C3FF2A8713FD4D47C3FF2A8713FD4D47C3F8BEF6D3EEEC77C3F8BEF6D3EEEC77C3F9F6131BFEDBA7C3F9F6131BFEDBA7C3FBB297BBFD3AD7C3FBB297BBFD3AD7C3F9C0DBCBE9FA07C3F9C0DBCBE9FA07C3FB68E153F52937C3FB68E153F52937C3F99A37F3FEA857C3F99A37F3FEA857C3FE45FFD3E68787C3FE45FFD3E68787C3FFA7AEDBECD6A7C3FFA7AEDBECD6A7C3F9DFF7EBF185D7C3F9DFF7EBF185D7C3FF6CF1CBF494F7C3FF6CF1CBF494F7C3FBA17AB3E60417C3FBA17AB3E60417C3F0E41793F5D337C3F0E41793F5D337C3F7FCC373F40257C3F7FCC373F40257C3FA48F4ABE0A177C3FA48F4ABE0A177C3F5B856EBFBA087C3F5B856EBFBA087C3F461B4FBF50FA7B3F461B4FBF50FA7B3FC5876B3DCCEB7B3FC5876B3DCCEB7B3F84035F3F2EDD7B3F84035F3F2EDD7B3FDF44623F77CE7B3FDF44623F77CE7B3FEDF2AB3DA6BF7B3FEDF2AB3DA6BF7B3FFD0A4BBFBBB07B3FFD0A4BBFBBB07B3F1AE770BFB6A17B3F1AE770BFB6A17B3F591C65BE97927B3F591C65BE97927B3F1A02333F5F837B3F1A02333F5F837B3FFCB67A3F0D747B3FFCB67A3F0D747B3FA7D4B73EA1647B3FA7D4B73EA1647B3F016417BF1C557B3F016417BF1C557B3F3E827FBF7C457B3F3E827FBF7C457B3F366DF9BEC3357B3F366DF9BEC3357B3F697CF13EF1257B3F697CF13EF1257B3F51307F3F04167B3F51307F3F04167B3FE0031B3FFE057B3FE0031B3FFE057B3F7A5BAFBEDEF57A3F7A5BAFBEDEF57A3FD8C279BFA5E57A3FD8C279BFA5E57A3FDF3636BF51D57A3FDF3636BF51D57A3F1470533EE5C47A3F1470533EE5C47A3FA2556F3F5EB47A3FA2556F3F5EB47A3F3BC44D3FBEA37A3F3BC44D3FBEA37A3FA9DB87BD04937A3FA9DB87BD04937A3F1C1E60BF30827A3F1C1E60BF30827A3F463361BF43717A3F463361BF43717A3FF3E099BD3C607A3FF3E099BD3C607A3F3F6A4C3F1B4F7A3F3F6A4C3F1B4F7A3F6D20703FE13D7A3F6D20703FE13D7A3F8E445C3E8D2C7A3F8E445C3E8D2C7A3FFE9E34BF201B7A3FFE9E34BF201B7A3F353F7ABF99097A3F353F7ABF99097A3F0298B3BEF8F7793F0298B3BEF8F7793F4336193F3EE6793F4336193F3EE6793FC35B7F3F6AD4793FC35B7F3F6AD4793F8875F53E7DC2793F8875F53E7DC2793FFF78F5BE76B0793FFF78F5BE76B0793FE75B7FBF559E793FE75B7FBF559E793FAE3419BF1B8C793FAE3419BF1B8C793FB69BB33EC879793FB69BB33EC879793F9F3F7A3F5A67793F9F3F7A3F5A67793F979D343FD454793F979D343FD454793F464C5CBE3342793F464C5CBE3342793F1C2170BF7A2F793F1C2170BF7A2F793F0F694CBFA61C793F0F694CBFA61C793FB6F0993DB909793FB6F0993DB909793F3634613FB3F6783F3634613FB3F6783F281D603F93E3783F281D603F93E3783FE4CB873D5AD0783FE4CB873D5AD0783F67C54DBF07BD783F67C54DBF07BD783FEF546FBF9BA9783FEF546FBF9BA9783F586853BE1596783F586853BE1596783F4238363F7682783F4238363F7682783F69C2793FBE6E783F69C2793FBE6E783FC357AF3EEB5A783FC357AF3EEB5A783F72051BBF0047783F72051BBF0047783F29307FBFFB32783F29307FBFFB32783FED78F1BEDD1E783FED78F1BEDD1E783FAA70F93EA50A783FAA70F93EA50A783F5D827F3F54F6773F5D827F3F54F6773F6962173FE9E1773F6962173FE9E1773F57D8B7BE65CD773F57D8B7BE65CD773F62B77ABFC8B8773F62B77ABFC8B8773FB00033BF11A4773FB00033BF11A4773F0D24653E418F773F0D24653E418F773FC5E7703F587A773FC5E7703F587A773FC9094B3F5565773FC9094B3F5565773FAC02ACBD3950773FAC02ACBD3950773FCB4562BF043B773FCB4562BF043B773F8C025FBFB525773F8C025FBFB525773F36686BBD4D10773F36686BBD4D10773F6F1C4F3FCCFA763F6F1C4F3FCCFA763FA4846E3F31E5763FA4846E3F31E5763FE5874A3E7ECF763FE5874A3E7ECF763FDFCD37BFB1B9763FDFCD37BFB1B9763F9A4079BFCAA3763F9A4079BFCAA3763F0014ABBECA8D763F0014ABBECA8D763F86D11C3FB277763F86D11C3FB277763F70FF7E3F7F61763F70FF7E3F7F61763F7A77ED3E344B763F7A77ED3E344B763F5363FDBECF34763F5363FDBECF34763FB4A37FBF521E763FB4A37FBF521E763F1B8D15BFBB07763F1B8D15BFBB07763F4911BC3E0AF1753F4911BC3E0AF1753F1D2A7B3F41DA753F1D2A7B3F41DA753F3260313F5FC3753F3260313F5FC3753F3BF76DBE63AC753F3BF76DBE63AC753F99A971BF4E95753F99A971BF4E95753F71A649BF207E753F71A649BF207E753F3011BE3DD966753F3011BE3DD966753FD752633F784F753FD752633F784F753F76E35D3FFF37753F76E35D3FFF37753FEC33473D6D20753FEC33473D6D20753F506F50BFC108753F506F50BFC108753F90AF6DBFFCF0743F90AF6DBFFCF0743F62A341BE1ED9743F62A341BE1ED9743FCC5F393F28C1743FCC5F393F28C1743FCCB9783F18A9743FCCB9783F18A9743FCFCCA63EEF90743FCFCCA63EEF90743F749A1EBFAD78743F749A1EBFAD78743F9AC97EBF5260743F9AC97EBF5260743F4471E9BEDD47743F4471E9BEDD47743F74A8003F502F743F74A8003F502F743FEABF7F3FAA16743FEABF7F3FAA16743FCDB4133FEBFD733FCDB4133FEBFD733F7546C0BE13E5733F7546C0BE13E5733FCF977BBF22CC733FCF977BBF22CC733F25BC2FBF18B3733F25BC2FBF18B3733FA3C5763EF599733FA3C5763EF599733F9466723FB980733F9466723FB980733F0E3F483F6467733F0E3F483F6467733FE31BD0BDF64D733FE31BD0BDF64D733F535B64BF7034733F535B64BF7034733FEEBF5CBFD01A733FEEBF5CBFD01A733F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F07D7A33AA46A573F07D7A33AB7C7683FFFD6233BB7C7683FFFD6233BC381103E6AC2753BC381103E6AC2753BCFBD41BFDED6A33BCFBD41BFDED6A33B107C75BF75CCCC3B107C75BF75CCCC3B8C0F8FBEF8C1F53B8C0F8FBEF8C1F53B4630283FB15B0F3C4630283FB15B0F3C95467D3F57D6233C95467D3F57D6233C3201D33EED50383C3201D33EED50383CF8440BBF6FCB4C3CF8440BBF6FCB4C3C5CFF7FBFDD45613C5CFF7FBFDD45613CD85C09BF33C0753CD85C09BF33C0753C2220D73E381D853C2220D73E381D853C72987D3F495A8F3C72987D3F495A8F3C4479263F4C97993C4479263F4C97993C106893BE3ED4A33C106893BE3ED4A33C251E76BF2011AE3C251E76BF2011AE3CB34040BFF04DB83CB34040BFF04DB83C6979193EAD8AC23C6979193EAD8AC23CC8B6693F57C7CC3CC8B6693F57C7CC3C102F563FEB03D73C102F563FEB03D73C150511BC6940E13C150511BC6940E13CE7A158BFD07CEB3CE7A158BFD07CEB3CFCD367BF20B9F53CFCD367BF20B9F53C388707BE55F5FF3C388707BE55F5FF3C0837433FB918053D0837433FB918053D0ED5743FB9360A3D0ED5743FB9360A3D2AB48A3EAB540F3D2AB48A3EAB540F3DE9E329BF8E72143DE9E329BF8E72143DA4EF7CBF6290193DA4EF7CBF6290193D06DECEBE27AE1E3D06DECEBE27AE1E3D4D2A0D3FDBCB233D4D2A0D3FDBCB233D39FA7F3F7FE9283D39FA7F3F7FE9283DF771073F11072E3DF771073F11072E3DC23ADBBE9124333DC23ADBBE9124333D38E57DBFFF41383D38E57DBFFF41383DEABE24BF5A5F3D3DEABE24BF5A5F3D3D9FBD973EA27C423D9FBD973EA27C423D4BBB763FD699473D4BBB763FD699473DBCBF3E3FF5B64C3DBCBF3E3FF5B64C3DFB6D22BE00D4513DFB6D22BE00D4513D28A16ABFF4F0563D28A16ABFF4F0563D31EF54BFD30D5C3D31EF54BFD30D5C3DA103913C9C2A613DA103913C9C2A613DD0D4593F4D47663DD0D4593F4D47663D9ADB663FE7636B3D9ADB663FE7636B3DE813FD3D6880703DE813FD3D6880703D56AC44BFD19C753D56AC44BFD19C753D222974BF21B97A3D222974BF21B97A3DFF5586BE57D57F3DFF5586BE57D57F3D24942B3FBA78823D24942B3FBA78823DA0937C3FBA06853DA0937C3FBA06853DB4B6CA3EAD94873DB4B6CA3EAD94873DCD0C0FBF92228A3DCD0C0FBF92228A3DF4EF7FBF69B08C3DF4EF7FBF69B08C3D5E8405BF323E8F3D5E8405BF323E8F3DFD50DF3EECCB913DFD50DF3EECCB913DE72C7E3F9759943DE72C7E3F9759943D4301233F33E7963D4301233F33E7963D23109CBEBF74993D23109CBEBF74993D7F5377BF3C029C3D7F5377BF3C029C3DF23A3DBFA88F9E3DF23A3DBFA88F9E3D4B5F2B3E051DA13D4B5F2B3E051DA13DD4866B3F51AAA33DD4866B3F51AAA33D0BAB533F8C37A63D0BAB533F8C37A63DCE81D9BCB6C4A83DCE81D9BCB6C4A83D5B035BBFCF51AB3D5B035BBFCF51AB3D97DE65BFD6DEAD3D97DE65BFD6DEAD3D4D14EBBDCC6BB03D4D14EBBDCC6BB03DB31D463FAFF8B23DB31D463FAFF8B23D5178733F8085B53D5178733F8085B53D23F5813E3F12B83D23F5813E3F12B83DED402DBFEA9EBA3DED402DBFEA9EBA3D8C327CBF832BBD3D8C327CBF832BBD3D518BC6BE08B8BF3D518BC6BE08B8BF3D6EEC103F7A44C23D6EEC103F7A44C23D8EE07F3FD7D0C43D8EE07F3FD7D0C43D1894033F215DC73D1894033F215DC73DBC62E3BE56E9C93DBC62E3BE56E9C93D7D6F7EBF7675CC3D7D6F7EBF7675CC3D574021BF8201CF3D574021BF8201CF3D865FA03E788DD13D865FA03E788DD13DBCE6773F5919D43DBCE6773F5919D43D5CB23B3F24A5D63D5CB23B3F24A5D63D2B4D34BED930D93D2B4D34BED930D93DC6676CBF78BCDB3DC6676CBF78BCDB3DA76252BF0048DE3DA76252BF0048DE3DD0FD103D72D3E03DD0FD103D72D3E03D822D5C3FCD5EE33D822D5C3FCD5EE33DF7DC643F10EAE53DF7DC643F10EAE53DFB0FD93D3C75E83DFB0FD93D3C75E83D178B47BF5000EB3D178B47BF5000EB3D9EC272BF4C8BED3D9EC272BF4C8BED3D56237BBE2F16F03D56237BBE2F16F03D3DEA2E3FFAA0F23D3DEA2E3FFAA0F23D68CC7B3FAC2BF53D68CC7B3FAC2BF53DF25BC23E46B6F73DF25BC23E46B6F73D28C912BFC540FA3D28C912BFC540FA3D05CC7FBF2BCBFC3D05CC7FBF2BCBFC3D2EA101BF7755FF3D2EA101BF7755FF3DED6FE73ED5EF003EED6FE73ED5EF003EF8AC7E3FE134023EF8AC7E3FE134023E2F7C1F3FDF79033E2F7C1F3FDF79033EB1ABA4BED0BE043EB1ABA4BED0BE043E007578BFB303063E007578BFB303063E03263ABF8948073E03263ABF8948073E6D373D3E518D083E6D373D3E518D083EFA436D3F0BD2093EFA436D3F0BD2093E0B16513FB7160B3E0B16513FB7160B3ECF3735BD545B0C3ECF3735BD545B0C3E3E535DBFE39F0D3E3E535DBFE39F0D3EC0D663BF64E40E3EC0D663BF64E40E3E4F07C7BDD628103E4F07C7BDD628103E7AF4483F396D113E7AF4483F396D113E0D08723F8DB1123E0D08723F8DB1123E5D57723ED3F5133E5D57723ED3F5133E0B9030BF093A153E0B9030BF093A153E38617BBF307E163E38617BBF307E163EAE28BEBE47C2173EAE28BEBE47C2173EEFA2143F4F06193EEFA2143F4F06193E5BB27F3F484A1A3E5BB27F3F484A1A3E5557FF3E308E1B3E5557FF3E308E1B3E7878EBBE09D21C3E7878EBBE09D21C3E58E57EBFD1151E3E58E57EBFD1151E3ED4B41DBF89591F3ED4B41DBF89591F3E8FF4A83E319D203E8FF4A83E319D203E49FE783FC9E0213E49FE783FC9E0213EED95383F5024233EED95383F5024233EE41D46BEC667243EE41D46BEC667243E6C1B6EBF2CAB253E6C1B6EBF2CAB253E3DC54FBF80EE263E3DC54FBF80EE263E2D6E593DC331283E2D6E593DC331283E8A745E3FF574293E8A745E3FF574293EF8CB623F16B82A3EF8CB623F16B82A3EA4FAB43D26FB2B3EA4FAB43D26FB2B3ED5594ABF233E2D3ED5594ABF233E2D3EA14871BF0F812E3EA14871BF0F812E3E888669BEE9C32F3E888669BEE9C32F3E4F32323FB106313E4F32323FB106313EFCF07A3F6749323EFCF07A3F6749323E99F1B93E0B8C333E99F1B93E0B8C333EBC7916BF9CCE343EBC7916BF9CCE343E90937FBF1B11363E90937FBF1B11363E2F67FBBE8753373E2F67FBBE8753373E4B7CEF3EE195383E4B7CEF3EE195383E9B187F3F27D8393E9B187F3F27D8393E4FEA1B3F5B1A3B3E4FEA1B3F5B1A3B3E093AADBE7B5C3C3E093AADBE7B5C3C3E938279BF889E3D3E938279BF889E3D3E240237BF81E03E3E240237BF81E03E3E61004F3E6822403E61004F3E6822403E18EE6E3F3A64413E18EE6E3F3A64413E45704E3FF8A5423E45704E3FF8A5423E2EA07DBDA3E7433E2EA07DBDA3E7433E60915FBF3A29453E60915FBF3A29453EA3BC61BFBC6A463EA3BC61BFBC6A463E59EAA2BD2AAC473E59EAA2BD2AAC473E21BB4B3F84ED483E21BB4B3F84ED483E5D84703FC92E4A3E5D84703FC92E4A3E04B1603EF96F4B3E04B1603EF96F4B3EFFD033BF15B14C3EFFD033BF15B14C3EB97B7ABF1BF24D3EB97B7ABF1BF24D3EC9B6B5BE0D334F3EC9B6B5BE0D334F3E844D183FE973503E844D183FE973503EA46F7F3FB0B4513EA46F7F3FB0B4513EFE71F73E61F5523EFE71F73E61F5523E507BF3BEFD35543E507BF3BEFD35543EC0467FBF8376553EC0467FBF8376553EAA1C1ABFF3B6563EAA1C1ABFF3B6563E0A7CB13E4EF7573E0A7CB13E4EF7573EDC017A3F9237593EDC017A3F9237593EAF6A353FC0775A3EAF6A353FC0775A3EB7DE57BED7B75B3EB7DE57BED7B75B3EF9BB6FBFD8F75C3EF9BB6FBFD8F75C3E28174DBFC2375E3E28174DBFC2375E3E8CE6903D96775F3E8CE6903D96775F3EB9A9603F53B7603EB9A9603F53B7603EC7A8603FF9F6613EC7A8603FF9F6613EC9D6903D8736633EC9D6903D8736633E57184DBFFF75643E57184DBFFF75643E47BB6FBF5EB5653E47BB6FBF5EB5653EFED657BEA7F4663EFED657BEA7F4663E146C353FD833683E146C353FD833683E6F017A3FF172693E6F017A3FF172693E5578B13EF2B16A3E5578B13EF2B16A3E3E1E1ABFDBF06B3E3E1E1ABFDBF06B3E9A467FBFAC2F6D3E9A467FBFAC2F6D3ED777F3BE646E6E3ED777F3BE646E6E3E7375F73E04AD6F3E7375F73E04AD6F3EC66F7F3F8CEB703EC66F7F3F8CEB703EED4B183FFB29723EED4B183FFB29723E7BBAB5BE5168733E7BBAB5BE5168733E217C7ABF8EA6743E217C7ABF8EA6743E97CF33BFB3E4753E97CF33BFB3E4753EB9B8603EBE22773EB9B8603EBE22773E0B85703FAF60783E0B85703FAF60783EEFB94B3F889E793EEFB94B3F889E793E1AFAA2BD46DC7A3E1AFAA2BD46DC7A3E92BD61BFEB197C3E92BD61BFEB197C3E69905FBF77577D3E69905FBF77577D3EA2807DBDE8947E3EA2807DBDE8947E3E70714E3F3FD27F3E70714E3F3FD27F3E62ED6E3FBE87803E62ED6E3FBE87803EA4F84E3E4F26813EA4F84E3E4F26813E860337BFD4C4813E860337BFD4C4813E228279BF4A63823E228279BF4A63823E5136ADBEB401833E5136ADBEB401833EE0EB1B3F10A0833EE0EB1B3F10A0833E70187F3F5E3E843E70187F3F5E3E843ECD78EF3EA0DC843ECD78EF3EA0DC843EA06AFBBED37A853EA06AFBBED37A853EAD937FBFF918863EAD937FBFF918863E237816BF11B7863E237816BF11B7863E47F5B93E1B55873E47F5B93E1B55873E61F17A3F17F3873E61F17A3F17F3873EE430323F0691883EE430323F0691883E3A8E69BEE62E893E3A8E69BEE62E893E4A4971BFB9CC893E4A4971BFB9CC893E9F584ABF7D6A8A3E9F584ABF7D6A8A3E620AB53D33088B3E620AB53D33088B3EE3CC623FDBA58B3EE3CC623FDBA58B3E90735E3F75438C3E90735E3F75438C3E9D4E593D00E18C3E9D4E593D00E18C3E65C64FBF7D7E8D3E65C64FBF7D7E8D3EB31A6EBFEB1B8E3EB31A6EBFEB1B8E3E231646BE4BB98E3E231646BE4BB98E3E4C97383F9C568F3E4C97383F9C568F3ED3FD783FDFF38F3ED3FD783FDFF38F3ED4F0A83E1291903ED4F0A83E1291903E62B61DBF372E913E62B61DBF372E913E29E57EBF4DCB913E29E57EBF4DCB913EF674EBBE5468923EF674EBBE5468923EC15AFF3E4D05933EC15AFF3E4D05933E73B27F3F36A2933E73B27F3F36A2933E53A1143F103F943E53A1143F103F943E592CBEBEDBDB943E592CBEBEDBDB943E97617BBF9678953E97617BBF9678953E9D8E30BF4215963E9D8E30BF4215963E0B5F723EDFB1963E0B5F723EDFB1963EB208723F6D4E973EB208723F6D4E973E40F3483FEBEA973E40F3483FEBEA973E0917C7BD5987983E0917C7BD5987983EA7D763BFB823993EA7D763BFB823993E40525DBF07C0993E40525DBF07C0993E3C1835BD465C9A3E3C1835BD465C9A3E2F17513F76F89A3E2F17513F76F89A3E3C436D3F96949B3E3C436D3F96949B3EA92F3D3EA5309C3EA92F3D3EA5309C3E5E273ABFA5CC9C3E5E273ABFA5CC9C3E867478BF95689D3E867478BF95689D3EF3A7A4BE74049E3EF3A7A4BE74049E3EBA7D1F3F44A09E3EBA7D1F3F44A09E3EC5AC7E3F033C9F3EC5AC7E3F033C9F3E666CE73EB2D79F3E666CE73EB2D79F3EE2A201BF5073A03EE2A201BF5073A03E"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x3727C5AC> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.reshape %arg27 : (tensor<256xi32>) -> tensor<256x1xi32> + %13 = "stablehlo.gather"(%arg0, %12) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<10x8xf32>, tensor<256x1xi32>) -> tensor<256x8xf32> + %14 = stablehlo.reshape %arg28 : (tensor<256xi32>) -> tensor<256x1xi32> + %15 = "stablehlo.gather"(%0, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %16 = "stablehlo.gather"(%1, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %17 = stablehlo.iota dim = 0 : tensor<256xi32> + %18 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<256xi32>) -> tensor<256x256xi32> + %19 = stablehlo.iota dim = 0 : tensor<256xi32> + %20 = stablehlo.broadcast_in_dim %19, dims = [1] : (tensor<256xi32>) -> tensor<256x256xi32> + %21 = stablehlo.compare LE, %20, %18, SIGNED : (tensor<256x256xi32>, tensor<256x256xi32>) -> tensor<256x256xi1> + %22 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x256xf32> + %23 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256x256xf32> + %24 = stablehlo.select %21, %22, %23 : tensor<256x256xi1>, tensor<256x256xf32> + %25 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %26 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %27 = stablehlo.multiply %13, %13 : tensor<256x8xf32> + %28 = stablehlo.reduce(%27 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %29 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %30 = stablehlo.divide %28, %29 : tensor<256xf32> + %31 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %32 = stablehlo.add %30, %31 : tensor<256xf32> + %33 = stablehlo.rsqrt %32 : tensor<256xf32> + %34 = stablehlo.broadcast_in_dim %33, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %35 = stablehlo.multiply %13, %34 : tensor<256x8xf32> + %36 = stablehlo.broadcast_in_dim %arg5, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %37 = stablehlo.multiply %35, %36 : tensor<256x8xf32> + %38 = stablehlo.dot_general %37, %arg10, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %39 = stablehlo.broadcast_in_dim %arg13, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %40 = stablehlo.add %38, %39 : tensor<256x8xf32> + %41 = stablehlo.reshape %40 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %42 = stablehlo.dot_general %37, %arg8, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %43 = stablehlo.broadcast_in_dim %arg12, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %44 = stablehlo.add %42, %43 : tensor<256x4xf32> + %45 = stablehlo.reshape %44 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %46 = stablehlo.dot_general %37, %arg11, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %47 = stablehlo.broadcast_in_dim %arg14, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %48 = stablehlo.add %46, %47 : tensor<256x4xf32> + %49 = stablehlo.reshape %48 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %50 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %51 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %52 = stablehlo.multiply %41, %50 : tensor<256x2x4xf32> + %53 = stablehlo.slice %41 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %54 = stablehlo.slice %41 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %55 = stablehlo.negate %54 : tensor<256x2x2xf32> + %56 = stablehlo.concatenate %55, %53, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %57 = stablehlo.multiply %56, %51 : tensor<256x2x4xf32> + %58 = stablehlo.add %52, %57 : tensor<256x2x4xf32> + %59 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %60 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %61 = stablehlo.multiply %45, %59 : tensor<256x1x4xf32> + %62 = stablehlo.slice %45 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %63 = stablehlo.slice %45 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %64 = stablehlo.negate %63 : tensor<256x1x2xf32> + %65 = stablehlo.concatenate %64, %62, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %66 = stablehlo.multiply %65, %60 : tensor<256x1x4xf32> + %67 = stablehlo.add %61, %66 : tensor<256x1x4xf32> + %68 = stablehlo.reshape %67 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %69 = stablehlo.dynamic_update_slice %25, %68, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %70 = stablehlo.reshape %49 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %71 = stablehlo.dynamic_update_slice %26, %70, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %72 = stablehlo.reshape %58 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %73 = stablehlo.dot_general %72, %67, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %74 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %75 = stablehlo.multiply %73, %74 : tensor<1x256x2x256xf32> + %76 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %77 = stablehlo.add %75, %76 : tensor<1x256x2x256xf32> + %78 = stablehlo.reduce(%77 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %79 = stablehlo.broadcast_in_dim %78, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %80 = stablehlo.subtract %77, %79 : tensor<1x256x2x256xf32> + %81 = stablehlo.exponential %80 : tensor<1x256x2x256xf32> + %82 = stablehlo.reduce(%81 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %83 = stablehlo.broadcast_in_dim %82, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %84 = stablehlo.divide %81, %83 : tensor<1x256x2x256xf32> + %85 = stablehlo.dot_general %84, %49, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %86 = stablehlo.transpose %85, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %87 = stablehlo.reshape %86 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %88 = stablehlo.dot_general %87, %arg9, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %89 = stablehlo.add %13, %88 : tensor<256x8xf32> + %90 = stablehlo.multiply %89, %89 : tensor<256x8xf32> + %91 = stablehlo.reduce(%90 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %92 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %93 = stablehlo.divide %91, %92 : tensor<256xf32> + %94 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %95 = stablehlo.add %93, %94 : tensor<256xf32> + %96 = stablehlo.rsqrt %95 : tensor<256xf32> + %97 = stablehlo.broadcast_in_dim %96, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %98 = stablehlo.multiply %89, %97 : tensor<256x8xf32> + %99 = stablehlo.broadcast_in_dim %arg6, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %100 = stablehlo.multiply %98, %99 : tensor<256x8xf32> + %101 = stablehlo.dot_general %100, %arg4, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %102 = stablehlo.dot_general %100, %arg7, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %103 = stablehlo.negate %101 : tensor<256x16xf32> + %104 = stablehlo.exponential %103 : tensor<256x16xf32> + %105 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %106 = stablehlo.add %105, %104 : tensor<256x16xf32> + %107 = stablehlo.divide %105, %106 : tensor<256x16xf32> + %108 = stablehlo.multiply %101, %107 : tensor<256x16xf32> + %109 = stablehlo.multiply %108, %102 : tensor<256x16xf32> + %110 = stablehlo.dot_general %109, %arg3, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %111 = stablehlo.add %89, %110 : tensor<256x8xf32> + %112 = stablehlo.multiply %111, %111 : tensor<256x8xf32> + %113 = stablehlo.reduce(%112 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %114 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %115 = stablehlo.divide %113, %114 : tensor<256xf32> + %116 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %117 = stablehlo.add %115, %116 : tensor<256xf32> + %118 = stablehlo.rsqrt %117 : tensor<256xf32> + %119 = stablehlo.broadcast_in_dim %118, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %120 = stablehlo.multiply %111, %119 : tensor<256x8xf32> + %121 = stablehlo.broadcast_in_dim %arg17, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %122 = stablehlo.multiply %120, %121 : tensor<256x8xf32> + %123 = stablehlo.dot_general %122, %arg22, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %124 = stablehlo.broadcast_in_dim %arg25, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %125 = stablehlo.add %123, %124 : tensor<256x8xf32> + %126 = stablehlo.reshape %125 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %127 = stablehlo.dot_general %122, %arg20, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %128 = stablehlo.broadcast_in_dim %arg24, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %129 = stablehlo.add %127, %128 : tensor<256x4xf32> + %130 = stablehlo.reshape %129 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %131 = stablehlo.dot_general %122, %arg23, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %132 = stablehlo.broadcast_in_dim %arg26, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %133 = stablehlo.add %131, %132 : tensor<256x4xf32> + %134 = stablehlo.reshape %133 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %135 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %136 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %137 = stablehlo.multiply %126, %135 : tensor<256x2x4xf32> + %138 = stablehlo.slice %126 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %139 = stablehlo.slice %126 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %140 = stablehlo.negate %139 : tensor<256x2x2xf32> + %141 = stablehlo.concatenate %140, %138, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %142 = stablehlo.multiply %141, %136 : tensor<256x2x4xf32> + %143 = stablehlo.add %137, %142 : tensor<256x2x4xf32> + %144 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %145 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %146 = stablehlo.multiply %130, %144 : tensor<256x1x4xf32> + %147 = stablehlo.slice %130 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %148 = stablehlo.slice %130 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %149 = stablehlo.negate %148 : tensor<256x1x2xf32> + %150 = stablehlo.concatenate %149, %147, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %151 = stablehlo.multiply %150, %145 : tensor<256x1x4xf32> + %152 = stablehlo.add %146, %151 : tensor<256x1x4xf32> + %153 = stablehlo.reshape %152 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %154 = stablehlo.dynamic_update_slice %69, %153, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %155 = stablehlo.reshape %134 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %156 = stablehlo.dynamic_update_slice %71, %155, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %157 = stablehlo.reshape %143 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %158 = stablehlo.dot_general %157, %152, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %159 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %160 = stablehlo.multiply %158, %159 : tensor<1x256x2x256xf32> + %161 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %162 = stablehlo.add %160, %161 : tensor<1x256x2x256xf32> + %163 = stablehlo.reduce(%162 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %164 = stablehlo.broadcast_in_dim %163, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %165 = stablehlo.subtract %162, %164 : tensor<1x256x2x256xf32> + %166 = stablehlo.exponential %165 : tensor<1x256x2x256xf32> + %167 = stablehlo.reduce(%166 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %168 = stablehlo.broadcast_in_dim %167, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %169 = stablehlo.divide %166, %168 : tensor<1x256x2x256xf32> + %170 = stablehlo.dot_general %169, %134, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %171 = stablehlo.transpose %170, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %172 = stablehlo.reshape %171 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %173 = stablehlo.dot_general %172, %arg21, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %174 = stablehlo.add %111, %173 : tensor<256x8xf32> + %175 = stablehlo.multiply %174, %174 : tensor<256x8xf32> + %176 = stablehlo.reduce(%175 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %177 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %178 = stablehlo.divide %176, %177 : tensor<256xf32> + %179 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %180 = stablehlo.add %178, %179 : tensor<256xf32> + %181 = stablehlo.rsqrt %180 : tensor<256xf32> + %182 = stablehlo.broadcast_in_dim %181, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %183 = stablehlo.multiply %174, %182 : tensor<256x8xf32> + %184 = stablehlo.broadcast_in_dim %arg18, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %185 = stablehlo.multiply %183, %184 : tensor<256x8xf32> + %186 = stablehlo.dot_general %185, %arg16, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %187 = stablehlo.dot_general %185, %arg19, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %188 = stablehlo.negate %186 : tensor<256x16xf32> + %189 = stablehlo.exponential %188 : tensor<256x16xf32> + %190 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %191 = stablehlo.add %190, %189 : tensor<256x16xf32> + %192 = stablehlo.divide %190, %191 : tensor<256x16xf32> + %193 = stablehlo.multiply %186, %192 : tensor<256x16xf32> + %194 = stablehlo.multiply %193, %187 : tensor<256x16xf32> + %195 = stablehlo.dot_general %194, %arg15, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %196 = stablehlo.add %174, %195 : tensor<256x8xf32> + %197 = stablehlo.multiply %196, %196 : tensor<256x8xf32> + %198 = stablehlo.reduce(%197 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %199 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %200 = stablehlo.divide %198, %199 : tensor<256xf32> + %201 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %202 = stablehlo.add %200, %201 : tensor<256xf32> + %203 = stablehlo.rsqrt %202 : tensor<256xf32> + %204 = stablehlo.broadcast_in_dim %203, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %205 = stablehlo.multiply %196, %204 : tensor<256x8xf32> + %206 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %207 = stablehlo.multiply %205, %206 : tensor<256x8xf32> + %208 = stablehlo.constant dense<1> : tensor + %209 = stablehlo.subtract %arg29, %208 : tensor + %210 = stablehlo.dynamic_slice %207, %209, %9, sizes = [1, 8] : (tensor<256x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %211 = stablehlo.reshape %210 : (tensor<1x8xf32>) -> tensor<8xf32> + %212 = stablehlo.dot_general %211, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %213 = stablehlo.constant dense<0> : tensor + %214 = stablehlo.constant dense<0xFF800000> : tensor + %215 = stablehlo.iota dim = 0 : tensor<10xi32> + %216:2 = stablehlo.reduce(%212 init: %214), (%215 init: %213) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %217 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %218 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %219 = stablehlo.or %217, %218 : tensor + %220 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %221 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %222 = stablehlo.and %220, %221 : tensor + %223 = stablehlo.or %219, %222 : tensor + %224 = stablehlo.select %219, %amv_l, %amv_r : tensor, tensor + %225 = stablehlo.select %223, %ami_l, %ami_r : tensor, tensor + stablehlo.return %224, %225 : tensor, tensor + } + return %216#1, %154, %156 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/seed_oss/README.md b/src/lib/mlxcel-xla/assets/seed_oss/README.md new file mode 100644 index 000000000..f7248332a --- /dev/null +++ b/src/lib/mlxcel-xla/assets/seed_oss/README.md @@ -0,0 +1,19 @@ +# Seed-OSS config fixture (#499 dense arch pack) + +`config.json` is a **small synthetic** Seed-OSS config (real architecture switches, +tiny dims) registered as a byte-exact structural fixture in `validation.rs`. Real +Seed-OSS is 36B / 64 layers, so its full graphs would be multi-MB goldens; the real +`config.json` parsing is covered separately by the `config::tests` parse assertion. + +Seed-OSS (`model_type` `seed_oss`) is the proven **Qwen2** forward with standard +names: a **q/k/v projection bias** (from `attention_bias = true`; `o_proj` and the +MLP have none, and `attention_out_bias = false`), **untied** embeddings, and +`rope_type = "default"` served as plain RoPE. Verified from the HF `SeedOssModel` +source to use the emitter's half-split RoPE, `cat(freqs, freqs)` table, +`head_dim^-0.5` scaling, RMSNorm and SwiGLU. + +Validation: `validation::dense_pack_families_reuse_proven_graphs` asserts it emits +byte-for-byte identical StableHLO to an untied Qwen2 reference, and +`spike/openxla/synthetic_arch_check.py` matches a tiny random `SeedOssForCausalLM`'s +last-token logits to ~1e-9 (forward parity). The full token-exact / serve gate runs +post-merge via `scripts/xla/validate_arch.sh` on the real checkpoint. diff --git a/src/lib/mlxcel-xla/assets/seed_oss/config.json b/src/lib/mlxcel-xla/assets/seed_oss/config.json new file mode 100644 index 000000000..650b93baa --- /dev/null +++ b/src/lib/mlxcel-xla/assets/seed_oss/config.json @@ -0,0 +1,17 @@ +{ + "model_type": "seed_oss", + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 16, + "num_hidden_layers": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 10000000.0, + "vocab_size": 10, + "tie_word_embeddings": false, + "attention_bias": true, + "attention_out_bias": false, + "rope_scaling": { "rope_type": "default" }, + "hidden_act": "silu" +} diff --git a/src/lib/mlxcel-xla/assets/seed_oss/decode.mlir b/src/lib/mlxcel-xla/assets/seed_oss/decode.mlir new file mode 100644 index 000000000..d78736d8b --- /dev/null +++ b/src/lib/mlxcel-xla/assets/seed_oss/decode.mlir @@ -0,0 +1,218 @@ +module @decode_step { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg4: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg5: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg6: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg7: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg10: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg11: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg12: tensor<4xf32> loc("params['layers'][0]['bk']"), %arg13: tensor<8xf32> loc("params['layers'][0]['bq']"), %arg14: tensor<4xf32> loc("params['layers'][0]['bv']"), %arg15: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg16: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg17: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg18: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg19: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg20: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg21: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg22: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg23: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg24: tensor<4xf32> loc("params['layers'][1]['bk']"), %arg25: tensor<8xf32> loc("params['layers'][1]['bq']"), %arg26: tensor<4xf32> loc("params['layers'][1]['bv']"), %arg27: tensor loc("token"), %arg28: tensor loc("pos"), %arg29: tensor loc("cache_len"), %arg30: tensor<2x256x1x4xf32> loc("kcache"), %arg31: tensor<2x256x1x4xf32> loc("vcache")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3FFFFF7F3F40510A3FFFFF7F3F3311D5BEFDFF7F3F3311D5BEFDFF7F3F26707DBFF8FF7F3F26707DBFF8FF7F3F305527BFF3FF7F3F305527BFF3FF7F3F2C3C913EEBFF7F3F2C3C913EEBFF7F3FB8CD753FE2FF7F3FB8CD753FE2FF7F3FBDFF403FD7FF7F3FBDFF403FD7FF7F3FF6FD14BECAFF7F3FF6FD14BECAFF7F3FD53F69BFBCFF7F3FD53F69BFBCFF7F3F64CD56BFACFF7F3F64CD56BFACFF7F3F7205913B9AFF7F3F7205913B9AFF7F3FD006583F87FF7F3FD006583F87FF7F3F6F4E683F72FF7F3F6F4E683F72FF7F3FD7040C3E5CFF7F3FD7040C3E5CFF7F3FE87A42BF43FF7F3FE87A42BF43FF7F3F2C2975BF29FF7F3F2C2975BF29FF7F3F36E28CBE0EFF7F3F36E28CBE0EFF7F3F840A293FF0FE7F3F840A293FF0FE7F3FBF1B7D3FD1FE7F3FBF1B7D3FD1FE7F3F22F0D03EB0FE7F3F22F0D03EB0FE7F3FFC370CBF8EFE7F3FFC370CBF8EFE7F3F6FFD7FBF6AFE7F3F6FFD7FBF6AFE7F3FBF6708BF44FE7F3FBF6708BF44FE7F3FFE2DD93E1DFE7F3FFE2DD93E1DFE7F3F78BF7D3FF4FD7F3F78BF7D3FF4FD7F3F819C253FC9FD7F3F819C253FC9FD7F3F389395BE9CFD7F3F389395BE9CFD7F3F576D76BF6EFD7F3F576D76BF6EFD7F3FB3803FBF3FFD7F3FB3803FBF3FFD7F3F18F41D3E0DFD7F3F18F41D3E0DFD7F3F8E2C6A3FDAFC7F3F8E2C6A3FDAFC7F3FAA8F553FA5FC7F3FAA8F553FA5FC7F3FB78659BC6EFC7F3FB78659BC6EFC7F3FE73B59BF36FC7F3FE73B59BF36FC7F3F5F5867BFFCFB7F3F5F5867BFFCFB7F3FEA0803BEC1FB7F3FEA0803BEC1FB7F3F2DF2433F84FB7F3F2DF2433F84FB7F3FB57F743F45FB7F3FB57F743F45FB7F3F6C85883E04FB7F3F6C85883E04FB7F3F74BC2ABFC2FA7F3F74BC2ABFC2FA7F3F44C27CBF7EFA7F3F44C27CBF7EFA7F3FE0CACCBE38FA7F3FE0CACCBE38FA7F3FE81B0E3FF1F97F3FE81B0E3FF1F97F3FBBF57F3FA8F97F3FBBF57F3FA8F97F3F807B063F5DF97F3F807B063F5DF97F3F6D46DDBE11F97F3F6D46DDBE11F97F3FB3097EBFC3F87F3FB3097EBFC3F87F3F80E023BF73F87F3F80E023BF73F87F3F44E7993E22F87F3F44E7993E22F87F3F0308773FCFF77F3F0308773FCFF77F3FD1FD3D3F7AF77F3FD1FD3D3F7AF77F3F0EE726BE24F77F3F0EE726BE24F77F3F95146BBFCCF67F3F95146BBFCCF67F3FA64D54BF72F67F3FA64D54BF72F67F3F2C43B53C17F67F3F2C43B53C17F67F3FA26C5A3FB9F57F3FA26C5A3FB9F57F3FAC5D663F5BF57F3FAC5D663F5BF57F3FB714F43DFAF47F3FB714F43DFAF47F3F836545BF98F47F3F836545BF98F47F3F56D173BF34F47F3F56D173BF34F47F3FE62584BECFF37F3FE62584BECFF37F3FF76A2C3F68F37F3FF76A2C3F68F37F3FB8637C3FFFF27F3FB8637C3FFFF27F3F83A1C83E94F27F3F83A1C83E94F27F3FFAFC0FBF28F27F3FFAFC0FBF28F27F3FE5E87FBFBAF17F3FE5E87FBFBAF17F3F908C04BF4AF17F3F908C04BF4AF17F3F6D5AE13ED9F07F3F6D5AE13ED9F07F3FD54E7E3F66F07F3FD54E7E3F66F07F3F3521223FF2EF7F3F3521223FF2EF7F3F3A389EBE7BEF7F3F3A389EBE7BEF7F3FBC9D77BF04EF7F3FBC9D77BF04EF7F3F20773CBF8AEE7F3F20773CBF8AEE7F3FACD62F3E0FEE7F3FACD62F3E0FEE7F3FE4F76B3F92ED7F3FE4F76B3F92ED7F3F6107533F13ED7F3F6107533F13ED7F3F5ABFFDBC93EC7F3F5ABFFDBC93EC7F3FFC985BBF11EC7F3FFC985BBF11EC7F3F5A5E65BF8DEB7F3F5A5E65BF8DEB7F3FB512E2BD08EB7F3FB512E2BD08EB7F3FE4D4463F81EA7F3FE4D4463F81EA7F3F141E733FF8E97F3F141E733FF8E97F3F72877F3E6DE97F3F72877F3E6DE97F3F05162EBFE1E87F3F05162EBFE1E87F3F1C007CBF54E87F3F1C007CBF54E87F3F1F74C4BEC4E77F3F1F74C4BEC4E77F3F28DB113F33E77F3F28DB113F33E77F3FEDD67F3FA0E67F3FEDD67F3FA0E67F3FF69A023F0CE67F3FF69A023F0CE67F3FE869E5BE76E57F3FE869E5BE76E57F3FDE8E7EBFDEE47F3FDE8E7EBFDEE47F3FAA5E20BF44E47F3FAA5E20BF44E47F3F0486A23EA9E37F3F0486A23EA9E37F3F7D2E783F0CE37F3F7D2E783F0CE37F3FA8EC3A3F6EE27F3FA8EC3A3F6EE27F3FC2C238BECEE17F3FC2C238BECEE17F3F78D66CBF2CE17F3F78D66CBF2CE17F3FE0BC51BF88E07F3FE0BC51BF88E07F3F381B233DE3DF7F3F381B233DE3DF7F3FEEC05C3F3CDF7F3FEEC05C3F3CDF7F3F6E5A643F94DE7F3F6E5A643F94DE7F3F2A0CD03DE9DD7F3F2A0CD03DE9DD7F3F494048BF3DDD7F3F494048BF3DDD7F3FF16572BF90DC7F3FF16572BF90DC7F3FF8BD76BEE0DB7F3FF8BD76BEE0DB7F3F95BD2F3F2FDB7F3F95BD2F3F2FDB7F3F71977B3F7DDA7F3F71977B3F7DDA7F3FCB42C03EC8D97F3FCB42C03EC8D97F3F6AB613BF12D97F3F6AB613BF12D97F3FD4BF7FBF5BD87F3FD4BF7FBF5BD87F3FBFA600BFA1D77F3FBFA600BFA1D77F3FC874E93EE6D67F3FC874E93EE6D67F3FCBC97E3F2AD67F3FCBC97E3F2AD67F3FE7981E3F6BD57F3FE7981E3F6BD57F3F8BD0A6BEABD47F3F8BD0A6BEABD47F3F44BA78BFEAD37F3F44BA78BFEAD37F3F6F5E39BF26D37F3F6F5E39BF26D37F3F25AB413E61D27F3F25AB413E61D27F3F4CB06D3F9AD17F3F4CB06D3F9AD17F3F2A6E503FD2D07F3F2A6E503FD2D07F3F7E5347BD08D07F3F7E5347BD08D07F3F73E45DBF3CCF7F3F73E45DBF3CCF7F3FEE5163BF6ECE7F3FEE5163BF6ECE7F3F7301BEBD9FCD7F3F7301BEBD9FCD7F3FA9A7493FCFCC7F3FA9A7493FCFCC7F3FF2A8713FFCCB7F3FF2A8713FFCCB7F3F8BEF6D3E28CB7F3F8BEF6D3E28CB7F3F9F6131BF52CA7F3F9F6131BF52CA7F3FBB297BBF7AC97F3FBB297BBF7AC97F3F9C0DBCBEA1C87F3F9C0DBCBEA1C87F3FB68E153FC6C77F3FB68E153FC6C77F3F99A37F3FEAC67F3F99A37F3FEAC67F3FE45FFD3E0CC67F3FE45FFD3E0CC67F3FFA7AEDBE2CC57F3FFA7AEDBE2CC57F3F9DFF7EBF4AC47F3F9DFF7EBF4AC47F3FF6CF1CBF67C37F3FF6CF1CBF67C37F3FBA17AB3E82C27F3FBA17AB3E82C27F3F0E41793F9BC17F3F0E41793F9BC17F3F7FCC373FB3C07F3F7FCC373FB3C07F3FA48F4ABEC9BF7F3FA48F4ABEC9BF7F3F5B856EBFDDBE7F3F5B856EBFDDBE7F3F461B4FBFF0BD7F3F461B4FBFF0BD7F3FC5876B3D01BD7F3FC5876B3D01BD7F3F84035F3F10BC7F3F84035F3F10BC7F3FDF44623F1EBB7F3FDF44623F1EBB7F3FEDF2AB3D2ABA7F3FEDF2AB3D2ABA7F3FFD0A4BBF34B97F3FFD0A4BBF34B97F3F1AE770BF3DB87F3F1AE770BF3DB87F3F591C65BE44B77F3F591C65BE44B77F3F1A02333F49B67F3F1A02333F49B67F3FFCB67A3F4DB57F3FFCB67A3F4DB57F3FA7D4B73E4FB47F3FA7D4B73E4FB47F3F016417BF4FB37F3F016417BF4FB37F3F3E827FBF4EB27F3F3E827FBF4EB27F3F366DF9BE4AB17F3F366DF9BE4AB17F3F697CF13E46B07F3F697CF13E46B07F3F51307F3F3FAF7F3F51307F3F3FAF7F3FE0031B3F37AE7F3FE0031B3F37AE7F3F7A5BAFBE2DAD7F3F7A5BAFBE2DAD7F3FD8C279BF22AC7F3FD8C279BF22AC7F3FDF3636BF15AB7F3FDF3636BF15AB7F3F1470533E06AA7F3F1470533E06AA7F3FA2556F3FF5A87F3FA2556F3FF5A87F3F3BC44D3FE3A77F3F3BC44D3FE3A77F3FA9DB87BDCFA67F3FA9DB87BDCFA67F3F1C1E60BFBAA57F3F1C1E60BFBAA57F3F463361BFA2A47F3F463361BFA2A47F3FF3E099BD8AA37F3FF3E099BD8AA37F3F3F6A4C3F6FA27F3F3F6A4C3F6FA27F3F6D20703F53A17F3F6D20703F53A17F3F8E445C3E35A07F3F8E445C3E35A07F3FFE9E34BF159F7F3FFE9E34BF159F7F3F353F7ABFF49D7F3F353F7ABFF49D7F3F0298B3BED19C7F3F0298B3BED19C7F3F4336193FAC9B7F3F4336193FAC9B7F3FC35B7F3F869A7F3FC35B7F3F869A7F3F8875F53E5E997F3F8875F53E5E997F3FFF78F5BE35987F3FFF78F5BE35987F3FE75B7FBF09977F3FE75B7FBF09977F3FAE3419BFDC957F3FAE3419BFDC957F3FB69BB33EAE947F3FB69BB33EAE947F3F9F3F7A3F7D937F3F9F3F7A3F7D937F3F979D343F4B927F3F979D343F4B927F3F464C5CBE18917F3F464C5CBE18917F3F1C2170BFE28F7F3F1C2170BFE28F7F3F0F694CBFAB8E7F3F0F694CBFAB8E7F3FB6F0993D728D7F3FB6F0993D728D7F3F3634613F388C7F3F3634613F388C7F3F281D603FFC8A7F3F281D603FFC8A7F3FE4CB873DBE897F3FE4CB873DBE897F3F67C54DBF7F887F3F67C54DBF7F887F3FEF546FBF3E877F3FEF546FBF3E877F3F586853BEFB857F3F586853BEFB857F3F4238363FB7847F3F4238363FB7847F3F69C2793F70837F3F69C2793F70837F3FC357AF3E29827F3FC357AF3E29827F3F72051BBFDF807F3F72051BBFDF807F3F29307FBF947F7F3F29307FBF947F7F3FED78F1BE477E7F3FED78F1BE477E7F3FAA70F93EF97C7F3FAA70F93EF97C7F3F5D827F3FA97B7F3F5D827F3FA97B7F3F6962173F577A7F3F6962173F577A7F3F57D8B7BE03797F3F57D8B7BE03797F3F62B77ABFAE777F3F62B77ABFAE777F3FB00033BF57767F3FB00033BF57767F3F0D24653EFF747F3F0D24653EFF747F3FC5E7703FA4737F3FC5E7703FA4737F3FC9094B3F49727F3FC9094B3F49727F3FAC02ACBDEB707F3FAC02ACBDEB707F3FCB4562BF8C6F7F3FCB4562BF8C6F7F3F8C025FBF2B6E7F3F8C025FBF2B6E7F3F36686BBDC86C7F3F36686BBDC86C7F3F6F1C4F3F646B7F3F6F1C4F3F646B7F3FA4846E3FFE697F3FA4846E3FFE697F3FE5874A3E97687F3FE5874A3E97687F3FDFCD37BF2D677F3FDFCD37BF2D677F3F9A4079BFC2657F3F9A4079BFC2657F3F0014ABBE56647F3F0014ABBE56647F3F86D11C3FE7627F3F86D11C3FE7627F3F70FF7E3F78617F3F70FF7E3F78617F3F7A77ED3E06607F3F7A77ED3E06607F3F5363FDBE935E7F3F5363FDBE935E7F3FB4A37FBF1E5D7F3FB4A37FBF1E5D7F3F1B8D15BFA75B7F3F1B8D15BFA75B7F3F4911BC3E2F5A7F3F4911BC3E2F5A7F3F1D2A7B3FB5587F3F1D2A7B3FB5587F3F3260313F39577F3F3260313F39577F3F3BF76DBEBC557F3F3BF76DBEBC557F3F99A971BF3D547F3F99A971BF3D547F3F71A649BFBC527F3F71A649BFBC527F3F3011BE3D39517F3F3011BE3D39517F3FD752633FB54F7F3FD752633FB54F7F3F76E35D3F304E7F3F76E35D3F304E7F3FEC33473DA84C7F3FEC33473DA84C7F3F506F50BF1F4B7F3F506F50BF1F4B7F3F90AF6DBF94497F3F90AF6DBF94497F3F62A341BE08487F3F62A341BE08487F3FCC5F393F7A467F3FCC5F393F7A467F3FCCB9783FEA447F3FCCB9783FEA447F3FCFCCA63E59437F3FCFCCA63E59437F3F749A1EBFC6417F3F749A1EBFC6417F3F9AC97EBF31407F3F9AC97EBF31407F3F4471E9BE9A3E7F3F4471E9BE9A3E7F3F74A8003F023D7F3F74A8003F023D7F3FEABF7F3F693B7F3FEABF7F3F693B7F3FCDB4133FCD397F3FCDB4133FCD397F3F7546C0BE30387F3F7546C0BE30387F3FCF977BBF91367F3FCF977BBF91367F3F25BC2FBFF1347F3F25BC2FBFF1347F3FA3C5763E4F337F3FA3C5763E4F337F3F9466723FAB317F3F9466723FAB317F3F0E3F483F05307F3F0E3F483F05307F3FE31BD0BD5E2E7F3FE31BD0BD5E2E7F3F535B64BFB52C7F3F535B64BFB52C7F3FEEBF5CBF0B2B7F3FEEBF5CBF0B2B7F3F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F5FCBA539A46A573F5FCBA539B7C7683F5FCB253AB7C7683F5FCB253AC381103E0DB1783AC381103E0DB1783ACFBD41BF5CCBA53ACFBD41BF5CCBA53A107C75BF313ECF3A107C75BF313ECF3A8C0F8FBE05B1F83A8C0F8FBE05B1F83A4630283FEC11113B4630283FEC11113B95467D3F54CB253B95467D3F54CB253B3201D33EBB843A3B3201D33EBB843A3BF8440BBF203E4F3BF8440BBF203E4F3B5CFF7FBF85F7633B5CFF7FBF85F7633BD85C09BFE8B0783BD85C09BFE8B0783B2220D73E25B5863B2220D73E25B5863B72987D3FD411913B72987D3FD411913B4479263F836E9B3B4479263F836E9B3B106893BE31CBA53B106893BE31CBA53B251E76BFDE27B03B251E76BFDE27B03BB34040BF8984BA3BB34040BF8984BA3B6979193E34E1C43B6979193E34E1C43BC8B6693FDD3DCF3BC8B6693FDD3DCF3B102F563F849AD93B102F563F849AD93B150511BC2BF7E33B150511BC2BF7E33BE7A158BFCF53EE3BE7A158BFCF53EE3BFCD367BF72B0F83BFCD367BF72B0F83B388707BE8A86013C388707BE8A86013C0837433FDAB4063C0837433FDAB4063C0ED5743F29E30B3C0ED5743F29E30B3C2AB48A3E7711113C2AB48A3E7711113CE9E329BFC43F163CE9E329BFC43F163CA4EF7CBF116E1B3CA4EF7CBF116E1B3C06DECEBE5C9C203C06DECEBE5C9C203C4D2A0D3FA6CA253C4D2A0D3FA6CA253C39FA7F3FEFF82A3C39FA7F3FEFF82A3CF771073F3727303CF771073F3727303CC23ADBBE7E55353CC23ADBBE7E55353C38E57DBFC3833A3C38E57DBFC3833A3CEABE24BF08B23F3CEABE24BF08B23F3C9FBD973E4BE0443C9FBD973E4BE0443C4BBB763F8C0E4A3C4BBB763F8C0E4A3CBCBF3E3FCD3C4F3CBCBF3E3FCD3C4F3CFB6D22BE0C6B543CFB6D22BE0C6B543C28A16ABF4A99593C28A16ABF4A99593C31EF54BF86C75E3C31EF54BF86C75E3CA103913CC1F5633CA103913CC1F5633CD0D4593FFA23693CD0D4593FFA23693C9ADB663F32526E3C9ADB663F32526E3CE813FD3D6880733CE813FD3D6880733C56AC44BF9DAE783C56AC44BF9DAE783C222974BFD0DC7D3C222974BFD0DC7D3CFF5586BE8185813CFF5586BE8185813C24942B3F991C843C24942B3F991C843CA0937C3FB0B3863CA0937C3FB0B3863CB4B6CA3EC64A893CB4B6CA3EC64A893CCD0C0FBFDBE18B3CCD0C0FBFDBE18B3CF4EF7FBFEF788E3CF4EF7FBFEF788E3C5E8405BF0210913C5E8405BF0210913CFD50DF3E15A7933CFD50DF3E15A7933CE72C7E3F263E963CE72C7E3F263E963C4301233F37D5983C4301233F37D5983C23109CBE466C9B3C23109CBE466C9B3C7F5377BF55039E3C7F5377BF55039E3CF23A3DBF629AA03CF23A3DBF629AA03C4B5F2B3E6E31A33C4B5F2B3E6E31A33CD4866B3F7AC8A53CD4866B3F7AC8A53C0BAB533F845FA83C0BAB533F845FA83CCE81D9BC8DF6AA3CCE81D9BC8DF6AA3C5B035BBF958DAD3C5B035BBF958DAD3C97DE65BF9C24B03C97DE65BF9C24B03C4D14EBBDA1BBB23C4D14EBBDA1BBB23CB31D463FA652B53CB31D463FA652B53C5178733FA9E9B73C5178733FA9E9B73C23F5813EAB80BA3C23F5813EAB80BA3CED402DBFAC17BD3CED402DBFAC17BD3C8C327CBFACAEBF3C8C327CBFACAEBF3C518BC6BEAA45C23C518BC6BEAA45C23C6EEC103FA7DCC43C6EEC103FA7DCC43C8EE07F3FA373C73C8EE07F3FA373C73C1894033F9D0ACA3C1894033F9D0ACA3CBC62E3BE97A1CC3CBC62E3BE97A1CC3C7D6F7EBF8E38CF3C7D6F7EBF8E38CF3C574021BF85CFD13C574021BF85CFD13C865FA03E7A66D43C865FA03E7A66D43CBCE6773F6EFDD63CBCE6773F6EFDD63C5CB23B3F6094D93C5CB23B3F6094D93C2B4D34BE512BDC3C2B4D34BE512BDC3CC6676CBF40C2DE3CC6676CBF40C2DE3CA76252BF2E59E13CA76252BF2E59E13CD0FD103D1BF0E33CD0FD103D1BF0E33C822D5C3F0687E63C822D5C3F0687E63CF7DC643FEF1DE93CF7DC643FEF1DE93CFB0FD93DD7B4EB3CFB0FD93DD7B4EB3C178B47BFBE4BEE3C178B47BFBE4BEE3C9EC272BFA3E2F03C9EC272BFA3E2F03C56237BBE8679F33C56237BBE8679F33C3DEA2E3F6810F63C3DEA2E3F6810F63C68CC7B3F48A7F83C68CC7B3F48A7F83CF25BC23E263EFB3CF25BC23E263EFB3C28C912BF03D5FD3C28C912BF03D5FD3C05CC7FBFEF35003D05CC7FBFEF35003D2EA101BF5C81013D2EA101BF5C81013DED6FE73EC8CC023DED6FE73EC8CC023DF8AC7E3F3318043DF8AC7E3F3318043D2F7C1F3F9D63053D2F7C1F3F9D63053DB1ABA4BE06AF063DB1ABA4BE06AF063D007578BF6EFA073D007578BF6EFA073D03263ABFD645093D03263ABFD645093D6D373D3E3D910A3D6D373D3E3D910A3DFA436D3FA2DC0B3DFA436D3FA2DC0B3D0B16513F07280D3D0B16513F07280D3DCF3735BD6B730E3DCF3735BD6B730E3D3E535DBFCEBE0F3D3E535DBFCEBE0F3DC0D663BF300A113DC0D663BF300A113D4F07C7BD9155123D4F07C7BD9155123D7AF4483FF1A0133D7AF4483FF1A0133D0D08723F50EC143D0D08723F50EC143D5D57723EAE37163D5D57723EAE37163D0B9030BF0B83173D0B9030BF0B83173D38617BBF67CE183D38617BBF67CE183DAE28BEBEC3191A3DAE28BEBEC3191A3DEFA2143F1D651B3DEFA2143F1D651B3D5BB27F3F76B01C3D5BB27F3F76B01C3D5557FF3ECEFB1D3D5557FF3ECEFB1D3D7878EBBE25471F3D7878EBBE25471F3D58E57EBF7B92203D58E57EBF7B92203DD4B41DBFD0DD213DD4B41DBFD0DD213D8FF4A83E2429233D8FF4A83E2429233D49FE783F7774243D49FE783F7774243DED95383FC9BF253DED95383FC9BF253DE41D46BE190B273DE41D46BE190B273D6C1B6EBF6956283D6C1B6EBF6956283D3DC54FBFB7A1293D3DC54FBFB7A1293D2D6E593D05ED2A3D2D6E593D05ED2A3D8A745E3F51382C3D8A745E3F51382C3DF8CB623F9C832D3DF8CB623F9C832D3DA4FAB43DE6CE2E3DA4FAB43DE6CE2E3DD5594ABF2F1A303DD5594ABF2F1A303DA14871BF7665313DA14871BF7665313D888669BEBDB0323D888669BEBDB0323D4F32323F02FC333D4F32323F02FC333DFCF07A3F4647353DFCF07A3F4647353D99F1B93E8992363D99F1B93E8992363DBC7916BFCBDD373DBC7916BFCBDD373D90937FBF0C29393D90937FBF0C29393D2F67FBBE4B743A3D2F67FBBE4B743A3D4B7CEF3E89BF3B3D4B7CEF3E89BF3B3D9B187F3FC60A3D3D9B187F3FC60A3D3D4FEA1B3F02563E3D4FEA1B3F02563E3D093AADBE3CA13F3D093AADBE3CA13F3D938279BF75EC403D938279BF75EC403D240237BFAD37423D240237BFAD37423D61004F3EE482433D61004F3EE482433D18EE6E3F19CE443D18EE6E3F19CE443D45704E3F4D19463D45704E3F4D19463D2EA07DBD8064473D2EA07DBD8064473D60915FBFB2AF483D60915FBFB2AF483DA3BC61BFE2FA493DA3BC61BFE2FA493D59EAA2BD11464B3D59EAA2BD11464B3D21BB4B3F3E914C3D21BB4B3F3E914C3D5D84703F6ADC4D3D5D84703F6ADC4D3D04B1603E95274F3D04B1603E95274F3DFFD033BFBE72503DFFD033BFBE72503DB97B7ABFE6BD513DB97B7ABFE6BD513DC9B6B5BE0D09533DC9B6B5BE0D09533D844D183F3254543D844D183F3254543DA46F7F3F569F553DA46F7F3F569F553DFE71F73E79EA563DFE71F73E79EA563D507BF3BE9A35583D507BF3BE9A35583DC0467FBFBA80593DC0467FBFBA80593DAA1C1ABFD8CB5A3DAA1C1ABFD8CB5A3D0A7CB13EF5165C3D0A7CB13EF5165C3DDC017A3F10625D3DDC017A3F10625D3DAF6A353F2AAD5E3DAF6A353F2AAD5E3DB7DE57BE43F85F3DB7DE57BE43F85F3DF9BB6FBF5A43613DF9BB6FBF5A43613D28174DBF6F8E623D28174DBF6F8E623D8CE6903D83D9633D8CE6903D83D9633DB9A9603F9624653DB9A9603F9624653DC7A8603FA76F663DC7A8603FA76F663DC9D6903DB6BA673DC9D6903DB6BA673D57184DBFC405693D57184DBFC405693D47BB6FBFD1506A3D47BB6FBFD1506A3DFED657BEDC9B6B3DFED657BEDC9B6B3D146C353FE5E66C3D146C353FE5E66C3D6F017A3FED316E3D6F017A3FED316E3D5578B13EF37C6F3D5578B13EF37C6F3D3E1E1ABFF8C7703D3E1E1ABFF8C7703D9A467FBFFB12723D9A467FBFFB12723DD777F3BEFD5D733DD777F3BEFD5D733D7375F73EFDA8743D7375F73EFDA8743DC66F7F3FFBF3753DC66F7F3FFBF3753DED4B183FF83E773DED4B183FF83E773D7BBAB5BEF389783D7BBAB5BEF389783D217C7ABFEDD4793D217C7ABFEDD4793D97CF33BFE51F7B3D97CF33BFE51F7B3DB9B8603EDB6A7C3DB9B8603EDB6A7C3D0B85703FCFB57D3D0B85703FCFB57D3DEFB94B3FC2007F3DEFB94B3FC2007F3D1AFAA2BDDA25803D1AFAA2BDDA25803D92BD61BF52CB803D92BD61BF52CB803D69905FBFC870813D69905FBFC870813DA2807DBD3F16823DA2807DBD3F16823D70714E3FB4BB823D70714E3FB4BB823D62ED6E3F2861833D62ED6E3F2861833DA4F84E3E9C06843DA4F84E3E9C06843D860337BF0EAC843D860337BF0EAC843D228279BF8051853D228279BF8051853D5136ADBEF1F6853D5136ADBEF1F6853DE0EB1B3F619C863DE0EB1B3F619C863D70187F3FD041873D70187F3FD041873DCD78EF3E3EE7873DCD78EF3E3EE7873DA06AFBBEAC8C883DA06AFBBEAC8C883DAD937FBF1832893DAD937FBF1832893D237816BF84D7893D237816BF84D7893D47F5B93EEF7C8A3D47F5B93EEF7C8A3D61F17A3F58228B3D61F17A3F58228B3DE430323FC1C78B3DE430323FC1C78B3D3A8E69BE296D8C3D3A8E69BE296D8C3D4A4971BF90128D3D4A4971BF90128D3D9F584ABFF6B78D3D9F584ABFF6B78D3D620AB53D5B5D8E3D620AB53D5B5D8E3DE3CC623FC0028F3DE3CC623FC0028F3D90735E3F23A88F3D90735E3F23A88F3D9D4E593D854D903D9D4E593D854D903D65C64FBFE7F2903D65C64FBFE7F2903DB31A6EBF4798913DB31A6EBF4798913D231646BEA73D923D231646BEA73D923D4C97383F05E3923D4C97383F05E3923DD3FD783F6388933DD3FD783F6388933DD4F0A83EBF2D943DD4F0A83EBF2D943D62B61DBF1BD3943D62B61DBF1BD3943D29E57EBF7678953D29E57EBF7678953DF674EBBECF1D963DF674EBBECF1D963DC15AFF3E28C3963DC15AFF3E28C3963D73B27F3F8068973D73B27F3F8068973D53A1143FD70D983D53A1143FD70D983D592CBEBE2CB3983D592CBEBE2CB3983D97617BBF8158993D97617BBF8158993D9D8E30BFD5FD993D9D8E30BFD5FD993D0B5F723E28A39A3D0B5F723E28A39A3DB208723F79489B3DB208723F79489B3D40F3483FCAED9B3D40F3483FCAED9B3D0917C7BD1A939C3D0917C7BD1A939C3DA7D763BF68389D3DA7D763BF68389D3D40525DBFB6DD9D3D40525DBFB6DD9D3D3C1835BD02839E3D3C1835BD02839E3D2F17513F4E289F3D2F17513F4E289F3D3C436D3F98CD9F3D3C436D3F98CD9F3DA92F3D3EE272A03DA92F3D3EE272A03D5E273ABF2A18A13D5E273ABF2A18A13D867478BF72BDA13D867478BF72BDA13DF3A7A4BEB862A23DF3A7A4BEB862A23DBA7D1F3FFD07A33DBA7D1F3FFD07A33DC5AC7E3F41ADA33DC5AC7E3F41ADA33D666CE73E8452A43D666CE73E8452A43DE2A201BFC6F7A43DE2A201BFC6F7A43D"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x358637BD> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.dynamic_slice %arg0, %arg27, %9, sizes = [1, 8] : (tensor<10x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %13 = stablehlo.reshape %12 : (tensor<1x8xf32>) -> tensor<8xf32> + %14 = stablehlo.dynamic_slice %0, %arg28, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %15 = stablehlo.reshape %14 : (tensor<1x4xf32>) -> tensor<4xf32> + %16 = stablehlo.dynamic_slice %1, %arg28, %9, sizes = [1, 4] : (tensor<256x4xf32>, tensor, tensor) -> tensor<1x4xf32> + %17 = stablehlo.reshape %16 : (tensor<1x4xf32>) -> tensor<4xf32> + %18 = stablehlo.iota dim = 0 : tensor<256xi32> + %19 = stablehlo.broadcast_in_dim %arg29, dims = [] : (tensor) -> tensor<256xi32> + %20 = stablehlo.compare LE, %18, %19, SIGNED : (tensor<256xi32>, tensor<256xi32>) -> tensor<256xi1> + %21 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256xf32> + %22 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256xf32> + %23 = stablehlo.select %20, %21, %22 : tensor<256xi1>, tensor<256xf32> + %24 = stablehlo.multiply %13, %13 : tensor<8xf32> + %25 = stablehlo.reduce(%24 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %26 = stablehlo.divide %25, %7 : tensor + %27 = stablehlo.add %26, %6 : tensor + %28 = stablehlo.rsqrt %27 : tensor + %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor) -> tensor<8xf32> + %30 = stablehlo.multiply %13, %29 : tensor<8xf32> + %31 = stablehlo.multiply %30, %arg5 : tensor<8xf32> + %32 = stablehlo.dot_general %31, %arg10, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %33 = stablehlo.add %32, %arg13 : tensor<8xf32> + %34 = stablehlo.reshape %33 : (tensor<8xf32>) -> tensor<2x4xf32> + %35 = stablehlo.dot_general %31, %arg8, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %36 = stablehlo.add %35, %arg12 : tensor<4xf32> + %37 = stablehlo.reshape %36 : (tensor<4xf32>) -> tensor<1x4xf32> + %38 = stablehlo.dot_general %31, %arg11, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %39 = stablehlo.add %38, %arg14 : tensor<4xf32> + %40 = stablehlo.reshape %39 : (tensor<4xf32>) -> tensor<1x4xf32> + %41 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %42 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %43 = stablehlo.multiply %34, %41 : tensor<2x4xf32> + %44 = stablehlo.slice %34 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %45 = stablehlo.slice %34 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %46 = stablehlo.negate %45 : tensor<2x2xf32> + %47 = stablehlo.concatenate %46, %44, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %48 = stablehlo.multiply %47, %42 : tensor<2x4xf32> + %49 = stablehlo.add %43, %48 : tensor<2x4xf32> + %50 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %51 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %52 = stablehlo.multiply %37, %50 : tensor<1x4xf32> + %53 = stablehlo.slice %37 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %54 = stablehlo.slice %37 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %55 = stablehlo.negate %54 : tensor<1x2xf32> + %56 = stablehlo.concatenate %55, %53, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %57 = stablehlo.multiply %56, %51 : tensor<1x4xf32> + %58 = stablehlo.add %52, %57 : tensor<1x4xf32> + %59 = stablehlo.reshape %58 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %60 = stablehlo.dynamic_update_slice %arg30, %59, %10, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %61 = stablehlo.reshape %40 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %62 = stablehlo.dynamic_update_slice %arg31, %61, %10, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %63 = stablehlo.slice %60 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %64 = stablehlo.reshape %63 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %65 = stablehlo.slice %62 [0:1, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %66 = stablehlo.reshape %65 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %67 = stablehlo.reshape %49 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %68 = stablehlo.dot_general %67, %64, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %69 = stablehlo.reshape %68 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %70 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %71 = stablehlo.multiply %69, %70 : tensor<2x256xf32> + %72 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %73 = stablehlo.add %71, %72 : tensor<2x256xf32> + %74 = stablehlo.reduce(%73 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %75 = stablehlo.broadcast_in_dim %74, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %76 = stablehlo.subtract %73, %75 : tensor<2x256xf32> + %77 = stablehlo.exponential %76 : tensor<2x256xf32> + %78 = stablehlo.reduce(%77 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %79 = stablehlo.broadcast_in_dim %78, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %80 = stablehlo.divide %77, %79 : tensor<2x256xf32> + %81 = stablehlo.reshape %80 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %82 = stablehlo.dot_general %81, %66, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %83 = stablehlo.reshape %82 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %84 = stablehlo.reshape %83 : (tensor<2x4xf32>) -> tensor<8xf32> + %85 = stablehlo.dot_general %84, %arg9, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %86 = stablehlo.add %13, %85 : tensor<8xf32> + %87 = stablehlo.multiply %86, %86 : tensor<8xf32> + %88 = stablehlo.reduce(%87 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %89 = stablehlo.divide %88, %7 : tensor + %90 = stablehlo.add %89, %6 : tensor + %91 = stablehlo.rsqrt %90 : tensor + %92 = stablehlo.broadcast_in_dim %91, dims = [] : (tensor) -> tensor<8xf32> + %93 = stablehlo.multiply %86, %92 : tensor<8xf32> + %94 = stablehlo.multiply %93, %arg6 : tensor<8xf32> + %95 = stablehlo.dot_general %94, %arg4, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %96 = stablehlo.dot_general %94, %arg7, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %97 = stablehlo.negate %95 : tensor<16xf32> + %98 = stablehlo.exponential %97 : tensor<16xf32> + %99 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %100 = stablehlo.add %99, %98 : tensor<16xf32> + %101 = stablehlo.divide %99, %100 : tensor<16xf32> + %102 = stablehlo.multiply %95, %101 : tensor<16xf32> + %103 = stablehlo.multiply %102, %96 : tensor<16xf32> + %104 = stablehlo.dot_general %103, %arg3, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %105 = stablehlo.add %86, %104 : tensor<8xf32> + %106 = stablehlo.multiply %105, %105 : tensor<8xf32> + %107 = stablehlo.reduce(%106 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %108 = stablehlo.divide %107, %7 : tensor + %109 = stablehlo.add %108, %6 : tensor + %110 = stablehlo.rsqrt %109 : tensor + %111 = stablehlo.broadcast_in_dim %110, dims = [] : (tensor) -> tensor<8xf32> + %112 = stablehlo.multiply %105, %111 : tensor<8xf32> + %113 = stablehlo.multiply %112, %arg17 : tensor<8xf32> + %114 = stablehlo.dot_general %113, %arg22, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %115 = stablehlo.add %114, %arg25 : tensor<8xf32> + %116 = stablehlo.reshape %115 : (tensor<8xf32>) -> tensor<2x4xf32> + %117 = stablehlo.dot_general %113, %arg20, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %118 = stablehlo.add %117, %arg24 : tensor<4xf32> + %119 = stablehlo.reshape %118 : (tensor<4xf32>) -> tensor<1x4xf32> + %120 = stablehlo.dot_general %113, %arg23, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<4x8xf32>) -> tensor<4xf32> + %121 = stablehlo.add %120, %arg26 : tensor<4xf32> + %122 = stablehlo.reshape %121 : (tensor<4xf32>) -> tensor<1x4xf32> + %123 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %124 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<2x4xf32> + %125 = stablehlo.multiply %116, %123 : tensor<2x4xf32> + %126 = stablehlo.slice %116 [0:2, 0:2] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %127 = stablehlo.slice %116 [0:2, 2:4] : (tensor<2x4xf32>) -> tensor<2x2xf32> + %128 = stablehlo.negate %127 : tensor<2x2xf32> + %129 = stablehlo.concatenate %128, %126, dim = 1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x4xf32> + %130 = stablehlo.multiply %129, %124 : tensor<2x4xf32> + %131 = stablehlo.add %125, %130 : tensor<2x4xf32> + %132 = stablehlo.broadcast_in_dim %15, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %133 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %134 = stablehlo.multiply %119, %132 : tensor<1x4xf32> + %135 = stablehlo.slice %119 [0:1, 0:2] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %136 = stablehlo.slice %119 [0:1, 2:4] : (tensor<1x4xf32>) -> tensor<1x2xf32> + %137 = stablehlo.negate %136 : tensor<1x2xf32> + %138 = stablehlo.concatenate %137, %135, dim = 1 : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32> + %139 = stablehlo.multiply %138, %133 : tensor<1x4xf32> + %140 = stablehlo.add %134, %139 : tensor<1x4xf32> + %141 = stablehlo.reshape %140 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %142 = stablehlo.dynamic_update_slice %60, %141, %11, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %143 = stablehlo.reshape %122 : (tensor<1x4xf32>) -> tensor<1x1x1x4xf32> + %144 = stablehlo.dynamic_update_slice %62, %143, %11, %arg29, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x1x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %145 = stablehlo.slice %142 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %146 = stablehlo.reshape %145 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %147 = stablehlo.slice %144 [1:2, 0:256, 0:1, 0:4] : (tensor<2x256x1x4xf32>) -> tensor<1x256x1x4xf32> + %148 = stablehlo.reshape %147 : (tensor<1x256x1x4xf32>) -> tensor<256x1x4xf32> + %149 = stablehlo.reshape %131 : (tensor<2x4xf32>) -> tensor<1x2x4xf32> + %150 = stablehlo.dot_general %149, %146, batching_dims = [0] x [1], contracting_dims = [2] x [2] : (tensor<1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x2x256xf32> + %151 = stablehlo.reshape %150 : (tensor<1x2x256xf32>) -> tensor<2x256xf32> + %152 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<2x256xf32> + %153 = stablehlo.multiply %151, %152 : tensor<2x256xf32> + %154 = stablehlo.broadcast_in_dim %23, dims = [1] : (tensor<256xf32>) -> tensor<2x256xf32> + %155 = stablehlo.add %153, %154 : tensor<2x256xf32> + %156 = stablehlo.reduce(%155 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %157 = stablehlo.broadcast_in_dim %156, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %158 = stablehlo.subtract %155, %157 : tensor<2x256xf32> + %159 = stablehlo.exponential %158 : tensor<2x256xf32> + %160 = stablehlo.reduce(%159 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<2x256xf32>, tensor) -> tensor<2xf32> + %161 = stablehlo.broadcast_in_dim %160, dims = [0] : (tensor<2xf32>) -> tensor<2x256xf32> + %162 = stablehlo.divide %159, %161 : tensor<2x256xf32> + %163 = stablehlo.reshape %162 : (tensor<2x256xf32>) -> tensor<1x2x256xf32> + %164 = stablehlo.dot_general %163, %148, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<1x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x2x4xf32> + %165 = stablehlo.reshape %164 : (tensor<1x2x4xf32>) -> tensor<2x4xf32> + %166 = stablehlo.reshape %165 : (tensor<2x4xf32>) -> tensor<8xf32> + %167 = stablehlo.dot_general %166, %arg21, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<8x8xf32>) -> tensor<8xf32> + %168 = stablehlo.add %105, %167 : tensor<8xf32> + %169 = stablehlo.multiply %168, %168 : tensor<8xf32> + %170 = stablehlo.reduce(%169 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %171 = stablehlo.divide %170, %7 : tensor + %172 = stablehlo.add %171, %6 : tensor + %173 = stablehlo.rsqrt %172 : tensor + %174 = stablehlo.broadcast_in_dim %173, dims = [] : (tensor) -> tensor<8xf32> + %175 = stablehlo.multiply %168, %174 : tensor<8xf32> + %176 = stablehlo.multiply %175, %arg18 : tensor<8xf32> + %177 = stablehlo.dot_general %176, %arg16, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %178 = stablehlo.dot_general %176, %arg19, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<16x8xf32>) -> tensor<16xf32> + %179 = stablehlo.negate %177 : tensor<16xf32> + %180 = stablehlo.exponential %179 : tensor<16xf32> + %181 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<16xf32> + %182 = stablehlo.add %181, %180 : tensor<16xf32> + %183 = stablehlo.divide %181, %182 : tensor<16xf32> + %184 = stablehlo.multiply %177, %183 : tensor<16xf32> + %185 = stablehlo.multiply %184, %178 : tensor<16xf32> + %186 = stablehlo.dot_general %185, %arg15, contracting_dims = [0] x [1] : (tensor<16xf32>, tensor<8x16xf32>) -> tensor<8xf32> + %187 = stablehlo.add %168, %186 : tensor<8xf32> + %188 = stablehlo.multiply %187, %187 : tensor<8xf32> + %189 = stablehlo.reduce(%188 init: %2) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor) -> tensor + %190 = stablehlo.divide %189, %7 : tensor + %191 = stablehlo.add %190, %6 : tensor + %192 = stablehlo.rsqrt %191 : tensor + %193 = stablehlo.broadcast_in_dim %192, dims = [] : (tensor) -> tensor<8xf32> + %194 = stablehlo.multiply %187, %193 : tensor<8xf32> + %195 = stablehlo.multiply %194, %arg1 : tensor<8xf32> + %196 = stablehlo.dot_general %195, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %197 = stablehlo.constant dense<0> : tensor + %198 = stablehlo.constant dense<0xFF800000> : tensor + %199 = stablehlo.iota dim = 0 : tensor<10xi32> + %200:2 = stablehlo.reduce(%196 init: %198), (%199 init: %197) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %201 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %202 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %203 = stablehlo.or %201, %202 : tensor + %204 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %205 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %206 = stablehlo.and %204, %205 : tensor + %207 = stablehlo.or %203, %206 : tensor + %208 = stablehlo.select %203, %amv_l, %amv_r : tensor, tensor + %209 = stablehlo.select %207, %ami_l, %ami_r : tensor, tensor + stablehlo.return %208, %209 : tensor, tensor + } + return %200#1, %142, %144 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/assets/seed_oss/prefill.mlir b/src/lib/mlxcel-xla/assets/seed_oss/prefill.mlir new file mode 100644 index 000000000..ba0f18dff --- /dev/null +++ b/src/lib/mlxcel-xla/assets/seed_oss/prefill.mlir @@ -0,0 +1,234 @@ +module @prefill { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8x16xf32> loc("params['layers'][0]['down']"), %arg4: tensor<16x8xf32> loc("params['layers'][0]['gate']"), %arg5: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg6: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg7: tensor<16x8xf32> loc("params['layers'][0]['up']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg9: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg10: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg11: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg12: tensor<4xf32> loc("params['layers'][0]['bk']"), %arg13: tensor<8xf32> loc("params['layers'][0]['bq']"), %arg14: tensor<4xf32> loc("params['layers'][0]['bv']"), %arg15: tensor<8x16xf32> loc("params['layers'][1]['down']"), %arg16: tensor<16x8xf32> loc("params['layers'][1]['gate']"), %arg17: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg18: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg19: tensor<16x8xf32> loc("params['layers'][1]['up']"), %arg20: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg21: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg22: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg23: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg24: tensor<4xf32> loc("params['layers'][1]['bk']"), %arg25: tensor<8xf32> loc("params['layers'][1]['bq']"), %arg26: tensor<4xf32> loc("params['layers'][1]['bv']"), %arg27: tensor<256xi32> loc("tokens"), %arg28: tensor<256xi32> loc("positions"), %arg29: tensor loc("real_len")) -> (tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3FFFFF7F3F40510A3FFFFF7F3F3311D5BEFDFF7F3F3311D5BEFDFF7F3F26707DBFF8FF7F3F26707DBFF8FF7F3F305527BFF3FF7F3F305527BFF3FF7F3F2C3C913EEBFF7F3F2C3C913EEBFF7F3FB8CD753FE2FF7F3FB8CD753FE2FF7F3FBDFF403FD7FF7F3FBDFF403FD7FF7F3FF6FD14BECAFF7F3FF6FD14BECAFF7F3FD53F69BFBCFF7F3FD53F69BFBCFF7F3F64CD56BFACFF7F3F64CD56BFACFF7F3F7205913B9AFF7F3F7205913B9AFF7F3FD006583F87FF7F3FD006583F87FF7F3F6F4E683F72FF7F3F6F4E683F72FF7F3FD7040C3E5CFF7F3FD7040C3E5CFF7F3FE87A42BF43FF7F3FE87A42BF43FF7F3F2C2975BF29FF7F3F2C2975BF29FF7F3F36E28CBE0EFF7F3F36E28CBE0EFF7F3F840A293FF0FE7F3F840A293FF0FE7F3FBF1B7D3FD1FE7F3FBF1B7D3FD1FE7F3F22F0D03EB0FE7F3F22F0D03EB0FE7F3FFC370CBF8EFE7F3FFC370CBF8EFE7F3F6FFD7FBF6AFE7F3F6FFD7FBF6AFE7F3FBF6708BF44FE7F3FBF6708BF44FE7F3FFE2DD93E1DFE7F3FFE2DD93E1DFE7F3F78BF7D3FF4FD7F3F78BF7D3FF4FD7F3F819C253FC9FD7F3F819C253FC9FD7F3F389395BE9CFD7F3F389395BE9CFD7F3F576D76BF6EFD7F3F576D76BF6EFD7F3FB3803FBF3FFD7F3FB3803FBF3FFD7F3F18F41D3E0DFD7F3F18F41D3E0DFD7F3F8E2C6A3FDAFC7F3F8E2C6A3FDAFC7F3FAA8F553FA5FC7F3FAA8F553FA5FC7F3FB78659BC6EFC7F3FB78659BC6EFC7F3FE73B59BF36FC7F3FE73B59BF36FC7F3F5F5867BFFCFB7F3F5F5867BFFCFB7F3FEA0803BEC1FB7F3FEA0803BEC1FB7F3F2DF2433F84FB7F3F2DF2433F84FB7F3FB57F743F45FB7F3FB57F743F45FB7F3F6C85883E04FB7F3F6C85883E04FB7F3F74BC2ABFC2FA7F3F74BC2ABFC2FA7F3F44C27CBF7EFA7F3F44C27CBF7EFA7F3FE0CACCBE38FA7F3FE0CACCBE38FA7F3FE81B0E3FF1F97F3FE81B0E3FF1F97F3FBBF57F3FA8F97F3FBBF57F3FA8F97F3F807B063F5DF97F3F807B063F5DF97F3F6D46DDBE11F97F3F6D46DDBE11F97F3FB3097EBFC3F87F3FB3097EBFC3F87F3F80E023BF73F87F3F80E023BF73F87F3F44E7993E22F87F3F44E7993E22F87F3F0308773FCFF77F3F0308773FCFF77F3FD1FD3D3F7AF77F3FD1FD3D3F7AF77F3F0EE726BE24F77F3F0EE726BE24F77F3F95146BBFCCF67F3F95146BBFCCF67F3FA64D54BF72F67F3FA64D54BF72F67F3F2C43B53C17F67F3F2C43B53C17F67F3FA26C5A3FB9F57F3FA26C5A3FB9F57F3FAC5D663F5BF57F3FAC5D663F5BF57F3FB714F43DFAF47F3FB714F43DFAF47F3F836545BF98F47F3F836545BF98F47F3F56D173BF34F47F3F56D173BF34F47F3FE62584BECFF37F3FE62584BECFF37F3FF76A2C3F68F37F3FF76A2C3F68F37F3FB8637C3FFFF27F3FB8637C3FFFF27F3F83A1C83E94F27F3F83A1C83E94F27F3FFAFC0FBF28F27F3FFAFC0FBF28F27F3FE5E87FBFBAF17F3FE5E87FBFBAF17F3F908C04BF4AF17F3F908C04BF4AF17F3F6D5AE13ED9F07F3F6D5AE13ED9F07F3FD54E7E3F66F07F3FD54E7E3F66F07F3F3521223FF2EF7F3F3521223FF2EF7F3F3A389EBE7BEF7F3F3A389EBE7BEF7F3FBC9D77BF04EF7F3FBC9D77BF04EF7F3F20773CBF8AEE7F3F20773CBF8AEE7F3FACD62F3E0FEE7F3FACD62F3E0FEE7F3FE4F76B3F92ED7F3FE4F76B3F92ED7F3F6107533F13ED7F3F6107533F13ED7F3F5ABFFDBC93EC7F3F5ABFFDBC93EC7F3FFC985BBF11EC7F3FFC985BBF11EC7F3F5A5E65BF8DEB7F3F5A5E65BF8DEB7F3FB512E2BD08EB7F3FB512E2BD08EB7F3FE4D4463F81EA7F3FE4D4463F81EA7F3F141E733FF8E97F3F141E733FF8E97F3F72877F3E6DE97F3F72877F3E6DE97F3F05162EBFE1E87F3F05162EBFE1E87F3F1C007CBF54E87F3F1C007CBF54E87F3F1F74C4BEC4E77F3F1F74C4BEC4E77F3F28DB113F33E77F3F28DB113F33E77F3FEDD67F3FA0E67F3FEDD67F3FA0E67F3FF69A023F0CE67F3FF69A023F0CE67F3FE869E5BE76E57F3FE869E5BE76E57F3FDE8E7EBFDEE47F3FDE8E7EBFDEE47F3FAA5E20BF44E47F3FAA5E20BF44E47F3F0486A23EA9E37F3F0486A23EA9E37F3F7D2E783F0CE37F3F7D2E783F0CE37F3FA8EC3A3F6EE27F3FA8EC3A3F6EE27F3FC2C238BECEE17F3FC2C238BECEE17F3F78D66CBF2CE17F3F78D66CBF2CE17F3FE0BC51BF88E07F3FE0BC51BF88E07F3F381B233DE3DF7F3F381B233DE3DF7F3FEEC05C3F3CDF7F3FEEC05C3F3CDF7F3F6E5A643F94DE7F3F6E5A643F94DE7F3F2A0CD03DE9DD7F3F2A0CD03DE9DD7F3F494048BF3DDD7F3F494048BF3DDD7F3FF16572BF90DC7F3FF16572BF90DC7F3FF8BD76BEE0DB7F3FF8BD76BEE0DB7F3F95BD2F3F2FDB7F3F95BD2F3F2FDB7F3F71977B3F7DDA7F3F71977B3F7DDA7F3FCB42C03EC8D97F3FCB42C03EC8D97F3F6AB613BF12D97F3F6AB613BF12D97F3FD4BF7FBF5BD87F3FD4BF7FBF5BD87F3FBFA600BFA1D77F3FBFA600BFA1D77F3FC874E93EE6D67F3FC874E93EE6D67F3FCBC97E3F2AD67F3FCBC97E3F2AD67F3FE7981E3F6BD57F3FE7981E3F6BD57F3F8BD0A6BEABD47F3F8BD0A6BEABD47F3F44BA78BFEAD37F3F44BA78BFEAD37F3F6F5E39BF26D37F3F6F5E39BF26D37F3F25AB413E61D27F3F25AB413E61D27F3F4CB06D3F9AD17F3F4CB06D3F9AD17F3F2A6E503FD2D07F3F2A6E503FD2D07F3F7E5347BD08D07F3F7E5347BD08D07F3F73E45DBF3CCF7F3F73E45DBF3CCF7F3FEE5163BF6ECE7F3FEE5163BF6ECE7F3F7301BEBD9FCD7F3F7301BEBD9FCD7F3FA9A7493FCFCC7F3FA9A7493FCFCC7F3FF2A8713FFCCB7F3FF2A8713FFCCB7F3F8BEF6D3E28CB7F3F8BEF6D3E28CB7F3F9F6131BF52CA7F3F9F6131BF52CA7F3FBB297BBF7AC97F3FBB297BBF7AC97F3F9C0DBCBEA1C87F3F9C0DBCBEA1C87F3FB68E153FC6C77F3FB68E153FC6C77F3F99A37F3FEAC67F3F99A37F3FEAC67F3FE45FFD3E0CC67F3FE45FFD3E0CC67F3FFA7AEDBE2CC57F3FFA7AEDBE2CC57F3F9DFF7EBF4AC47F3F9DFF7EBF4AC47F3FF6CF1CBF67C37F3FF6CF1CBF67C37F3FBA17AB3E82C27F3FBA17AB3E82C27F3F0E41793F9BC17F3F0E41793F9BC17F3F7FCC373FB3C07F3F7FCC373FB3C07F3FA48F4ABEC9BF7F3FA48F4ABEC9BF7F3F5B856EBFDDBE7F3F5B856EBFDDBE7F3F461B4FBFF0BD7F3F461B4FBFF0BD7F3FC5876B3D01BD7F3FC5876B3D01BD7F3F84035F3F10BC7F3F84035F3F10BC7F3FDF44623F1EBB7F3FDF44623F1EBB7F3FEDF2AB3D2ABA7F3FEDF2AB3D2ABA7F3FFD0A4BBF34B97F3FFD0A4BBF34B97F3F1AE770BF3DB87F3F1AE770BF3DB87F3F591C65BE44B77F3F591C65BE44B77F3F1A02333F49B67F3F1A02333F49B67F3FFCB67A3F4DB57F3FFCB67A3F4DB57F3FA7D4B73E4FB47F3FA7D4B73E4FB47F3F016417BF4FB37F3F016417BF4FB37F3F3E827FBF4EB27F3F3E827FBF4EB27F3F366DF9BE4AB17F3F366DF9BE4AB17F3F697CF13E46B07F3F697CF13E46B07F3F51307F3F3FAF7F3F51307F3F3FAF7F3FE0031B3F37AE7F3FE0031B3F37AE7F3F7A5BAFBE2DAD7F3F7A5BAFBE2DAD7F3FD8C279BF22AC7F3FD8C279BF22AC7F3FDF3636BF15AB7F3FDF3636BF15AB7F3F1470533E06AA7F3F1470533E06AA7F3FA2556F3FF5A87F3FA2556F3FF5A87F3F3BC44D3FE3A77F3F3BC44D3FE3A77F3FA9DB87BDCFA67F3FA9DB87BDCFA67F3F1C1E60BFBAA57F3F1C1E60BFBAA57F3F463361BFA2A47F3F463361BFA2A47F3FF3E099BD8AA37F3FF3E099BD8AA37F3F3F6A4C3F6FA27F3F3F6A4C3F6FA27F3F6D20703F53A17F3F6D20703F53A17F3F8E445C3E35A07F3F8E445C3E35A07F3FFE9E34BF159F7F3FFE9E34BF159F7F3F353F7ABFF49D7F3F353F7ABFF49D7F3F0298B3BED19C7F3F0298B3BED19C7F3F4336193FAC9B7F3F4336193FAC9B7F3FC35B7F3F869A7F3FC35B7F3F869A7F3F8875F53E5E997F3F8875F53E5E997F3FFF78F5BE35987F3FFF78F5BE35987F3FE75B7FBF09977F3FE75B7FBF09977F3FAE3419BFDC957F3FAE3419BFDC957F3FB69BB33EAE947F3FB69BB33EAE947F3F9F3F7A3F7D937F3F9F3F7A3F7D937F3F979D343F4B927F3F979D343F4B927F3F464C5CBE18917F3F464C5CBE18917F3F1C2170BFE28F7F3F1C2170BFE28F7F3F0F694CBFAB8E7F3F0F694CBFAB8E7F3FB6F0993D728D7F3FB6F0993D728D7F3F3634613F388C7F3F3634613F388C7F3F281D603FFC8A7F3F281D603FFC8A7F3FE4CB873DBE897F3FE4CB873DBE897F3F67C54DBF7F887F3F67C54DBF7F887F3FEF546FBF3E877F3FEF546FBF3E877F3F586853BEFB857F3F586853BEFB857F3F4238363FB7847F3F4238363FB7847F3F69C2793F70837F3F69C2793F70837F3FC357AF3E29827F3FC357AF3E29827F3F72051BBFDF807F3F72051BBFDF807F3F29307FBF947F7F3F29307FBF947F7F3FED78F1BE477E7F3FED78F1BE477E7F3FAA70F93EF97C7F3FAA70F93EF97C7F3F5D827F3FA97B7F3F5D827F3FA97B7F3F6962173F577A7F3F6962173F577A7F3F57D8B7BE03797F3F57D8B7BE03797F3F62B77ABFAE777F3F62B77ABFAE777F3FB00033BF57767F3FB00033BF57767F3F0D24653EFF747F3F0D24653EFF747F3FC5E7703FA4737F3FC5E7703FA4737F3FC9094B3F49727F3FC9094B3F49727F3FAC02ACBDEB707F3FAC02ACBDEB707F3FCB4562BF8C6F7F3FCB4562BF8C6F7F3F8C025FBF2B6E7F3F8C025FBF2B6E7F3F36686BBDC86C7F3F36686BBDC86C7F3F6F1C4F3F646B7F3F6F1C4F3F646B7F3FA4846E3FFE697F3FA4846E3FFE697F3FE5874A3E97687F3FE5874A3E97687F3FDFCD37BF2D677F3FDFCD37BF2D677F3F9A4079BFC2657F3F9A4079BFC2657F3F0014ABBE56647F3F0014ABBE56647F3F86D11C3FE7627F3F86D11C3FE7627F3F70FF7E3F78617F3F70FF7E3F78617F3F7A77ED3E06607F3F7A77ED3E06607F3F5363FDBE935E7F3F5363FDBE935E7F3FB4A37FBF1E5D7F3FB4A37FBF1E5D7F3F1B8D15BFA75B7F3F1B8D15BFA75B7F3F4911BC3E2F5A7F3F4911BC3E2F5A7F3F1D2A7B3FB5587F3F1D2A7B3FB5587F3F3260313F39577F3F3260313F39577F3F3BF76DBEBC557F3F3BF76DBEBC557F3F99A971BF3D547F3F99A971BF3D547F3F71A649BFBC527F3F71A649BFBC527F3F3011BE3D39517F3F3011BE3D39517F3FD752633FB54F7F3FD752633FB54F7F3F76E35D3F304E7F3F76E35D3F304E7F3FEC33473DA84C7F3FEC33473DA84C7F3F506F50BF1F4B7F3F506F50BF1F4B7F3F90AF6DBF94497F3F90AF6DBF94497F3F62A341BE08487F3F62A341BE08487F3FCC5F393F7A467F3FCC5F393F7A467F3FCCB9783FEA447F3FCCB9783FEA447F3FCFCCA63E59437F3FCFCCA63E59437F3F749A1EBFC6417F3F749A1EBFC6417F3F9AC97EBF31407F3F9AC97EBF31407F3F4471E9BE9A3E7F3F4471E9BE9A3E7F3F74A8003F023D7F3F74A8003F023D7F3FEABF7F3F693B7F3FEABF7F3F693B7F3FCDB4133FCD397F3FCDB4133FCD397F3F7546C0BE30387F3F7546C0BE30387F3FCF977BBF91367F3FCF977BBF91367F3F25BC2FBFF1347F3F25BC2FBFF1347F3FA3C5763E4F337F3FA3C5763E4F337F3F9466723FAB317F3F9466723FAB317F3F0E3F483F05307F3F0E3F483F05307F3FE31BD0BD5E2E7F3FE31BD0BD5E2E7F3F535B64BFB52C7F3F535B64BFB52C7F3FEEBF5CBF0B2B7F3FEEBF5CBF0B2B7F3F"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F5FCBA539A46A573F5FCBA539B7C7683F5FCB253AB7C7683F5FCB253AC381103E0DB1783AC381103E0DB1783ACFBD41BF5CCBA53ACFBD41BF5CCBA53A107C75BF313ECF3A107C75BF313ECF3A8C0F8FBE05B1F83A8C0F8FBE05B1F83A4630283FEC11113B4630283FEC11113B95467D3F54CB253B95467D3F54CB253B3201D33EBB843A3B3201D33EBB843A3BF8440BBF203E4F3BF8440BBF203E4F3B5CFF7FBF85F7633B5CFF7FBF85F7633BD85C09BFE8B0783BD85C09BFE8B0783B2220D73E25B5863B2220D73E25B5863B72987D3FD411913B72987D3FD411913B4479263F836E9B3B4479263F836E9B3B106893BE31CBA53B106893BE31CBA53B251E76BFDE27B03B251E76BFDE27B03BB34040BF8984BA3BB34040BF8984BA3B6979193E34E1C43B6979193E34E1C43BC8B6693FDD3DCF3BC8B6693FDD3DCF3B102F563F849AD93B102F563F849AD93B150511BC2BF7E33B150511BC2BF7E33BE7A158BFCF53EE3BE7A158BFCF53EE3BFCD367BF72B0F83BFCD367BF72B0F83B388707BE8A86013C388707BE8A86013C0837433FDAB4063C0837433FDAB4063C0ED5743F29E30B3C0ED5743F29E30B3C2AB48A3E7711113C2AB48A3E7711113CE9E329BFC43F163CE9E329BFC43F163CA4EF7CBF116E1B3CA4EF7CBF116E1B3C06DECEBE5C9C203C06DECEBE5C9C203C4D2A0D3FA6CA253C4D2A0D3FA6CA253C39FA7F3FEFF82A3C39FA7F3FEFF82A3CF771073F3727303CF771073F3727303CC23ADBBE7E55353CC23ADBBE7E55353C38E57DBFC3833A3C38E57DBFC3833A3CEABE24BF08B23F3CEABE24BF08B23F3C9FBD973E4BE0443C9FBD973E4BE0443C4BBB763F8C0E4A3C4BBB763F8C0E4A3CBCBF3E3FCD3C4F3CBCBF3E3FCD3C4F3CFB6D22BE0C6B543CFB6D22BE0C6B543C28A16ABF4A99593C28A16ABF4A99593C31EF54BF86C75E3C31EF54BF86C75E3CA103913CC1F5633CA103913CC1F5633CD0D4593FFA23693CD0D4593FFA23693C9ADB663F32526E3C9ADB663F32526E3CE813FD3D6880733CE813FD3D6880733C56AC44BF9DAE783C56AC44BF9DAE783C222974BFD0DC7D3C222974BFD0DC7D3CFF5586BE8185813CFF5586BE8185813C24942B3F991C843C24942B3F991C843CA0937C3FB0B3863CA0937C3FB0B3863CB4B6CA3EC64A893CB4B6CA3EC64A893CCD0C0FBFDBE18B3CCD0C0FBFDBE18B3CF4EF7FBFEF788E3CF4EF7FBFEF788E3C5E8405BF0210913C5E8405BF0210913CFD50DF3E15A7933CFD50DF3E15A7933CE72C7E3F263E963CE72C7E3F263E963C4301233F37D5983C4301233F37D5983C23109CBE466C9B3C23109CBE466C9B3C7F5377BF55039E3C7F5377BF55039E3CF23A3DBF629AA03CF23A3DBF629AA03C4B5F2B3E6E31A33C4B5F2B3E6E31A33CD4866B3F7AC8A53CD4866B3F7AC8A53C0BAB533F845FA83C0BAB533F845FA83CCE81D9BC8DF6AA3CCE81D9BC8DF6AA3C5B035BBF958DAD3C5B035BBF958DAD3C97DE65BF9C24B03C97DE65BF9C24B03C4D14EBBDA1BBB23C4D14EBBDA1BBB23CB31D463FA652B53CB31D463FA652B53C5178733FA9E9B73C5178733FA9E9B73C23F5813EAB80BA3C23F5813EAB80BA3CED402DBFAC17BD3CED402DBFAC17BD3C8C327CBFACAEBF3C8C327CBFACAEBF3C518BC6BEAA45C23C518BC6BEAA45C23C6EEC103FA7DCC43C6EEC103FA7DCC43C8EE07F3FA373C73C8EE07F3FA373C73C1894033F9D0ACA3C1894033F9D0ACA3CBC62E3BE97A1CC3CBC62E3BE97A1CC3C7D6F7EBF8E38CF3C7D6F7EBF8E38CF3C574021BF85CFD13C574021BF85CFD13C865FA03E7A66D43C865FA03E7A66D43CBCE6773F6EFDD63CBCE6773F6EFDD63C5CB23B3F6094D93C5CB23B3F6094D93C2B4D34BE512BDC3C2B4D34BE512BDC3CC6676CBF40C2DE3CC6676CBF40C2DE3CA76252BF2E59E13CA76252BF2E59E13CD0FD103D1BF0E33CD0FD103D1BF0E33C822D5C3F0687E63C822D5C3F0687E63CF7DC643FEF1DE93CF7DC643FEF1DE93CFB0FD93DD7B4EB3CFB0FD93DD7B4EB3C178B47BFBE4BEE3C178B47BFBE4BEE3C9EC272BFA3E2F03C9EC272BFA3E2F03C56237BBE8679F33C56237BBE8679F33C3DEA2E3F6810F63C3DEA2E3F6810F63C68CC7B3F48A7F83C68CC7B3F48A7F83CF25BC23E263EFB3CF25BC23E263EFB3C28C912BF03D5FD3C28C912BF03D5FD3C05CC7FBFEF35003D05CC7FBFEF35003D2EA101BF5C81013D2EA101BF5C81013DED6FE73EC8CC023DED6FE73EC8CC023DF8AC7E3F3318043DF8AC7E3F3318043D2F7C1F3F9D63053D2F7C1F3F9D63053DB1ABA4BE06AF063DB1ABA4BE06AF063D007578BF6EFA073D007578BF6EFA073D03263ABFD645093D03263ABFD645093D6D373D3E3D910A3D6D373D3E3D910A3DFA436D3FA2DC0B3DFA436D3FA2DC0B3D0B16513F07280D3D0B16513F07280D3DCF3735BD6B730E3DCF3735BD6B730E3D3E535DBFCEBE0F3D3E535DBFCEBE0F3DC0D663BF300A113DC0D663BF300A113D4F07C7BD9155123D4F07C7BD9155123D7AF4483FF1A0133D7AF4483FF1A0133D0D08723F50EC143D0D08723F50EC143D5D57723EAE37163D5D57723EAE37163D0B9030BF0B83173D0B9030BF0B83173D38617BBF67CE183D38617BBF67CE183DAE28BEBEC3191A3DAE28BEBEC3191A3DEFA2143F1D651B3DEFA2143F1D651B3D5BB27F3F76B01C3D5BB27F3F76B01C3D5557FF3ECEFB1D3D5557FF3ECEFB1D3D7878EBBE25471F3D7878EBBE25471F3D58E57EBF7B92203D58E57EBF7B92203DD4B41DBFD0DD213DD4B41DBFD0DD213D8FF4A83E2429233D8FF4A83E2429233D49FE783F7774243D49FE783F7774243DED95383FC9BF253DED95383FC9BF253DE41D46BE190B273DE41D46BE190B273D6C1B6EBF6956283D6C1B6EBF6956283D3DC54FBFB7A1293D3DC54FBFB7A1293D2D6E593D05ED2A3D2D6E593D05ED2A3D8A745E3F51382C3D8A745E3F51382C3DF8CB623F9C832D3DF8CB623F9C832D3DA4FAB43DE6CE2E3DA4FAB43DE6CE2E3DD5594ABF2F1A303DD5594ABF2F1A303DA14871BF7665313DA14871BF7665313D888669BEBDB0323D888669BEBDB0323D4F32323F02FC333D4F32323F02FC333DFCF07A3F4647353DFCF07A3F4647353D99F1B93E8992363D99F1B93E8992363DBC7916BFCBDD373DBC7916BFCBDD373D90937FBF0C29393D90937FBF0C29393D2F67FBBE4B743A3D2F67FBBE4B743A3D4B7CEF3E89BF3B3D4B7CEF3E89BF3B3D9B187F3FC60A3D3D9B187F3FC60A3D3D4FEA1B3F02563E3D4FEA1B3F02563E3D093AADBE3CA13F3D093AADBE3CA13F3D938279BF75EC403D938279BF75EC403D240237BFAD37423D240237BFAD37423D61004F3EE482433D61004F3EE482433D18EE6E3F19CE443D18EE6E3F19CE443D45704E3F4D19463D45704E3F4D19463D2EA07DBD8064473D2EA07DBD8064473D60915FBFB2AF483D60915FBFB2AF483DA3BC61BFE2FA493DA3BC61BFE2FA493D59EAA2BD11464B3D59EAA2BD11464B3D21BB4B3F3E914C3D21BB4B3F3E914C3D5D84703F6ADC4D3D5D84703F6ADC4D3D04B1603E95274F3D04B1603E95274F3DFFD033BFBE72503DFFD033BFBE72503DB97B7ABFE6BD513DB97B7ABFE6BD513DC9B6B5BE0D09533DC9B6B5BE0D09533D844D183F3254543D844D183F3254543DA46F7F3F569F553DA46F7F3F569F553DFE71F73E79EA563DFE71F73E79EA563D507BF3BE9A35583D507BF3BE9A35583DC0467FBFBA80593DC0467FBFBA80593DAA1C1ABFD8CB5A3DAA1C1ABFD8CB5A3D0A7CB13EF5165C3D0A7CB13EF5165C3DDC017A3F10625D3DDC017A3F10625D3DAF6A353F2AAD5E3DAF6A353F2AAD5E3DB7DE57BE43F85F3DB7DE57BE43F85F3DF9BB6FBF5A43613DF9BB6FBF5A43613D28174DBF6F8E623D28174DBF6F8E623D8CE6903D83D9633D8CE6903D83D9633DB9A9603F9624653DB9A9603F9624653DC7A8603FA76F663DC7A8603FA76F663DC9D6903DB6BA673DC9D6903DB6BA673D57184DBFC405693D57184DBFC405693D47BB6FBFD1506A3D47BB6FBFD1506A3DFED657BEDC9B6B3DFED657BEDC9B6B3D146C353FE5E66C3D146C353FE5E66C3D6F017A3FED316E3D6F017A3FED316E3D5578B13EF37C6F3D5578B13EF37C6F3D3E1E1ABFF8C7703D3E1E1ABFF8C7703D9A467FBFFB12723D9A467FBFFB12723DD777F3BEFD5D733DD777F3BEFD5D733D7375F73EFDA8743D7375F73EFDA8743DC66F7F3FFBF3753DC66F7F3FFBF3753DED4B183FF83E773DED4B183FF83E773D7BBAB5BEF389783D7BBAB5BEF389783D217C7ABFEDD4793D217C7ABFEDD4793D97CF33BFE51F7B3D97CF33BFE51F7B3DB9B8603EDB6A7C3DB9B8603EDB6A7C3D0B85703FCFB57D3D0B85703FCFB57D3DEFB94B3FC2007F3DEFB94B3FC2007F3D1AFAA2BDDA25803D1AFAA2BDDA25803D92BD61BF52CB803D92BD61BF52CB803D69905FBFC870813D69905FBFC870813DA2807DBD3F16823DA2807DBD3F16823D70714E3FB4BB823D70714E3FB4BB823D62ED6E3F2861833D62ED6E3F2861833DA4F84E3E9C06843DA4F84E3E9C06843D860337BF0EAC843D860337BF0EAC843D228279BF8051853D228279BF8051853D5136ADBEF1F6853D5136ADBEF1F6853DE0EB1B3F619C863DE0EB1B3F619C863D70187F3FD041873D70187F3FD041873DCD78EF3E3EE7873DCD78EF3E3EE7873DA06AFBBEAC8C883DA06AFBBEAC8C883DAD937FBF1832893DAD937FBF1832893D237816BF84D7893D237816BF84D7893D47F5B93EEF7C8A3D47F5B93EEF7C8A3D61F17A3F58228B3D61F17A3F58228B3DE430323FC1C78B3DE430323FC1C78B3D3A8E69BE296D8C3D3A8E69BE296D8C3D4A4971BF90128D3D4A4971BF90128D3D9F584ABFF6B78D3D9F584ABFF6B78D3D620AB53D5B5D8E3D620AB53D5B5D8E3DE3CC623FC0028F3DE3CC623FC0028F3D90735E3F23A88F3D90735E3F23A88F3D9D4E593D854D903D9D4E593D854D903D65C64FBFE7F2903D65C64FBFE7F2903DB31A6EBF4798913DB31A6EBF4798913D231646BEA73D923D231646BEA73D923D4C97383F05E3923D4C97383F05E3923DD3FD783F6388933DD3FD783F6388933DD4F0A83EBF2D943DD4F0A83EBF2D943D62B61DBF1BD3943D62B61DBF1BD3943D29E57EBF7678953D29E57EBF7678953DF674EBBECF1D963DF674EBBECF1D963DC15AFF3E28C3963DC15AFF3E28C3963D73B27F3F8068973D73B27F3F8068973D53A1143FD70D983D53A1143FD70D983D592CBEBE2CB3983D592CBEBE2CB3983D97617BBF8158993D97617BBF8158993D9D8E30BFD5FD993D9D8E30BFD5FD993D0B5F723E28A39A3D0B5F723E28A39A3DB208723F79489B3DB208723F79489B3D40F3483FCAED9B3D40F3483FCAED9B3D0917C7BD1A939C3D0917C7BD1A939C3DA7D763BF68389D3DA7D763BF68389D3D40525DBFB6DD9D3D40525DBFB6DD9D3D3C1835BD02839E3D3C1835BD02839E3D2F17513F4E289F3D2F17513F4E289F3D3C436D3F98CD9F3D3C436D3F98CD9F3DA92F3D3EE272A03DA92F3D3EE272A03D5E273ABF2A18A13D5E273ABF2A18A13D867478BF72BDA13D867478BF72BDA13DF3A7A4BEB862A23DF3A7A4BEB862A23DBA7D1F3FFD07A33DBA7D1F3FFD07A33DC5AC7E3F41ADA33DC5AC7E3F41ADA33D666CE73E8452A43D666CE73E8452A43DE2A201BFC6F7A43DE2A201BFC6F7A43D"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x358637BD> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.reshape %arg27 : (tensor<256xi32>) -> tensor<256x1xi32> + %13 = "stablehlo.gather"(%arg0, %12) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<10x8xf32>, tensor<256x1xi32>) -> tensor<256x8xf32> + %14 = stablehlo.reshape %arg28 : (tensor<256xi32>) -> tensor<256x1xi32> + %15 = "stablehlo.gather"(%0, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %16 = "stablehlo.gather"(%1, %14) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %17 = stablehlo.iota dim = 0 : tensor<256xi32> + %18 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<256xi32>) -> tensor<256x256xi32> + %19 = stablehlo.iota dim = 0 : tensor<256xi32> + %20 = stablehlo.broadcast_in_dim %19, dims = [1] : (tensor<256xi32>) -> tensor<256x256xi32> + %21 = stablehlo.compare LE, %20, %18, SIGNED : (tensor<256x256xi32>, tensor<256x256xi32>) -> tensor<256x256xi1> + %22 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x256xf32> + %23 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor) -> tensor<256x256xf32> + %24 = stablehlo.select %21, %22, %23 : tensor<256x256xi1>, tensor<256x256xf32> + %25 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %26 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %27 = stablehlo.multiply %13, %13 : tensor<256x8xf32> + %28 = stablehlo.reduce(%27 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %29 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %30 = stablehlo.divide %28, %29 : tensor<256xf32> + %31 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %32 = stablehlo.add %30, %31 : tensor<256xf32> + %33 = stablehlo.rsqrt %32 : tensor<256xf32> + %34 = stablehlo.broadcast_in_dim %33, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %35 = stablehlo.multiply %13, %34 : tensor<256x8xf32> + %36 = stablehlo.broadcast_in_dim %arg5, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %37 = stablehlo.multiply %35, %36 : tensor<256x8xf32> + %38 = stablehlo.dot_general %37, %arg10, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %39 = stablehlo.broadcast_in_dim %arg13, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %40 = stablehlo.add %38, %39 : tensor<256x8xf32> + %41 = stablehlo.reshape %40 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %42 = stablehlo.dot_general %37, %arg8, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %43 = stablehlo.broadcast_in_dim %arg12, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %44 = stablehlo.add %42, %43 : tensor<256x4xf32> + %45 = stablehlo.reshape %44 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %46 = stablehlo.dot_general %37, %arg11, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %47 = stablehlo.broadcast_in_dim %arg14, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %48 = stablehlo.add %46, %47 : tensor<256x4xf32> + %49 = stablehlo.reshape %48 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %50 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %51 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %52 = stablehlo.multiply %41, %50 : tensor<256x2x4xf32> + %53 = stablehlo.slice %41 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %54 = stablehlo.slice %41 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %55 = stablehlo.negate %54 : tensor<256x2x2xf32> + %56 = stablehlo.concatenate %55, %53, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %57 = stablehlo.multiply %56, %51 : tensor<256x2x4xf32> + %58 = stablehlo.add %52, %57 : tensor<256x2x4xf32> + %59 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %60 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %61 = stablehlo.multiply %45, %59 : tensor<256x1x4xf32> + %62 = stablehlo.slice %45 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %63 = stablehlo.slice %45 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %64 = stablehlo.negate %63 : tensor<256x1x2xf32> + %65 = stablehlo.concatenate %64, %62, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %66 = stablehlo.multiply %65, %60 : tensor<256x1x4xf32> + %67 = stablehlo.add %61, %66 : tensor<256x1x4xf32> + %68 = stablehlo.reshape %67 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %69 = stablehlo.dynamic_update_slice %25, %68, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %70 = stablehlo.reshape %49 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %71 = stablehlo.dynamic_update_slice %26, %70, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %72 = stablehlo.reshape %58 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %73 = stablehlo.dot_general %72, %67, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %74 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %75 = stablehlo.multiply %73, %74 : tensor<1x256x2x256xf32> + %76 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %77 = stablehlo.add %75, %76 : tensor<1x256x2x256xf32> + %78 = stablehlo.reduce(%77 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %79 = stablehlo.broadcast_in_dim %78, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %80 = stablehlo.subtract %77, %79 : tensor<1x256x2x256xf32> + %81 = stablehlo.exponential %80 : tensor<1x256x2x256xf32> + %82 = stablehlo.reduce(%81 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %83 = stablehlo.broadcast_in_dim %82, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %84 = stablehlo.divide %81, %83 : tensor<1x256x2x256xf32> + %85 = stablehlo.dot_general %84, %49, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %86 = stablehlo.transpose %85, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %87 = stablehlo.reshape %86 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %88 = stablehlo.dot_general %87, %arg9, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %89 = stablehlo.add %13, %88 : tensor<256x8xf32> + %90 = stablehlo.multiply %89, %89 : tensor<256x8xf32> + %91 = stablehlo.reduce(%90 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %92 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %93 = stablehlo.divide %91, %92 : tensor<256xf32> + %94 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %95 = stablehlo.add %93, %94 : tensor<256xf32> + %96 = stablehlo.rsqrt %95 : tensor<256xf32> + %97 = stablehlo.broadcast_in_dim %96, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %98 = stablehlo.multiply %89, %97 : tensor<256x8xf32> + %99 = stablehlo.broadcast_in_dim %arg6, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %100 = stablehlo.multiply %98, %99 : tensor<256x8xf32> + %101 = stablehlo.dot_general %100, %arg4, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %102 = stablehlo.dot_general %100, %arg7, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %103 = stablehlo.negate %101 : tensor<256x16xf32> + %104 = stablehlo.exponential %103 : tensor<256x16xf32> + %105 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %106 = stablehlo.add %105, %104 : tensor<256x16xf32> + %107 = stablehlo.divide %105, %106 : tensor<256x16xf32> + %108 = stablehlo.multiply %101, %107 : tensor<256x16xf32> + %109 = stablehlo.multiply %108, %102 : tensor<256x16xf32> + %110 = stablehlo.dot_general %109, %arg3, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %111 = stablehlo.add %89, %110 : tensor<256x8xf32> + %112 = stablehlo.multiply %111, %111 : tensor<256x8xf32> + %113 = stablehlo.reduce(%112 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %114 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %115 = stablehlo.divide %113, %114 : tensor<256xf32> + %116 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %117 = stablehlo.add %115, %116 : tensor<256xf32> + %118 = stablehlo.rsqrt %117 : tensor<256xf32> + %119 = stablehlo.broadcast_in_dim %118, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %120 = stablehlo.multiply %111, %119 : tensor<256x8xf32> + %121 = stablehlo.broadcast_in_dim %arg17, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %122 = stablehlo.multiply %120, %121 : tensor<256x8xf32> + %123 = stablehlo.dot_general %122, %arg22, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %124 = stablehlo.broadcast_in_dim %arg25, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %125 = stablehlo.add %123, %124 : tensor<256x8xf32> + %126 = stablehlo.reshape %125 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %127 = stablehlo.dot_general %122, %arg20, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %128 = stablehlo.broadcast_in_dim %arg24, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %129 = stablehlo.add %127, %128 : tensor<256x4xf32> + %130 = stablehlo.reshape %129 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %131 = stablehlo.dot_general %122, %arg23, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %132 = stablehlo.broadcast_in_dim %arg26, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %133 = stablehlo.add %131, %132 : tensor<256x4xf32> + %134 = stablehlo.reshape %133 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %135 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %136 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %137 = stablehlo.multiply %126, %135 : tensor<256x2x4xf32> + %138 = stablehlo.slice %126 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %139 = stablehlo.slice %126 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %140 = stablehlo.negate %139 : tensor<256x2x2xf32> + %141 = stablehlo.concatenate %140, %138, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %142 = stablehlo.multiply %141, %136 : tensor<256x2x4xf32> + %143 = stablehlo.add %137, %142 : tensor<256x2x4xf32> + %144 = stablehlo.broadcast_in_dim %15, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %145 = stablehlo.broadcast_in_dim %16, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %146 = stablehlo.multiply %130, %144 : tensor<256x1x4xf32> + %147 = stablehlo.slice %130 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %148 = stablehlo.slice %130 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %149 = stablehlo.negate %148 : tensor<256x1x2xf32> + %150 = stablehlo.concatenate %149, %147, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %151 = stablehlo.multiply %150, %145 : tensor<256x1x4xf32> + %152 = stablehlo.add %146, %151 : tensor<256x1x4xf32> + %153 = stablehlo.reshape %152 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %154 = stablehlo.dynamic_update_slice %69, %153, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %155 = stablehlo.reshape %134 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %156 = stablehlo.dynamic_update_slice %71, %155, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %157 = stablehlo.reshape %143 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %158 = stablehlo.dot_general %157, %152, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %159 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %160 = stablehlo.multiply %158, %159 : tensor<1x256x2x256xf32> + %161 = stablehlo.broadcast_in_dim %24, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %162 = stablehlo.add %160, %161 : tensor<1x256x2x256xf32> + %163 = stablehlo.reduce(%162 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %164 = stablehlo.broadcast_in_dim %163, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %165 = stablehlo.subtract %162, %164 : tensor<1x256x2x256xf32> + %166 = stablehlo.exponential %165 : tensor<1x256x2x256xf32> + %167 = stablehlo.reduce(%166 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %168 = stablehlo.broadcast_in_dim %167, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %169 = stablehlo.divide %166, %168 : tensor<1x256x2x256xf32> + %170 = stablehlo.dot_general %169, %134, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %171 = stablehlo.transpose %170, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %172 = stablehlo.reshape %171 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %173 = stablehlo.dot_general %172, %arg21, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %174 = stablehlo.add %111, %173 : tensor<256x8xf32> + %175 = stablehlo.multiply %174, %174 : tensor<256x8xf32> + %176 = stablehlo.reduce(%175 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %177 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %178 = stablehlo.divide %176, %177 : tensor<256xf32> + %179 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %180 = stablehlo.add %178, %179 : tensor<256xf32> + %181 = stablehlo.rsqrt %180 : tensor<256xf32> + %182 = stablehlo.broadcast_in_dim %181, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %183 = stablehlo.multiply %174, %182 : tensor<256x8xf32> + %184 = stablehlo.broadcast_in_dim %arg18, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %185 = stablehlo.multiply %183, %184 : tensor<256x8xf32> + %186 = stablehlo.dot_general %185, %arg16, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %187 = stablehlo.dot_general %185, %arg19, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<16x8xf32>) -> tensor<256x16xf32> + %188 = stablehlo.negate %186 : tensor<256x16xf32> + %189 = stablehlo.exponential %188 : tensor<256x16xf32> + %190 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x16xf32> + %191 = stablehlo.add %190, %189 : tensor<256x16xf32> + %192 = stablehlo.divide %190, %191 : tensor<256x16xf32> + %193 = stablehlo.multiply %186, %192 : tensor<256x16xf32> + %194 = stablehlo.multiply %193, %187 : tensor<256x16xf32> + %195 = stablehlo.dot_general %194, %arg15, contracting_dims = [1] x [1] : (tensor<256x16xf32>, tensor<8x16xf32>) -> tensor<256x8xf32> + %196 = stablehlo.add %174, %195 : tensor<256x8xf32> + %197 = stablehlo.multiply %196, %196 : tensor<256x8xf32> + %198 = stablehlo.reduce(%197 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %199 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %200 = stablehlo.divide %198, %199 : tensor<256xf32> + %201 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %202 = stablehlo.add %200, %201 : tensor<256xf32> + %203 = stablehlo.rsqrt %202 : tensor<256xf32> + %204 = stablehlo.broadcast_in_dim %203, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %205 = stablehlo.multiply %196, %204 : tensor<256x8xf32> + %206 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %207 = stablehlo.multiply %205, %206 : tensor<256x8xf32> + %208 = stablehlo.constant dense<1> : tensor + %209 = stablehlo.subtract %arg29, %208 : tensor + %210 = stablehlo.dynamic_slice %207, %209, %9, sizes = [1, 8] : (tensor<256x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %211 = stablehlo.reshape %210 : (tensor<1x8xf32>) -> tensor<8xf32> + %212 = stablehlo.dot_general %211, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + %213 = stablehlo.constant dense<0> : tensor + %214 = stablehlo.constant dense<0xFF800000> : tensor + %215 = stablehlo.iota dim = 0 : tensor<10xi32> + %216:2 = stablehlo.reduce(%212 init: %214), (%215 init: %213) across dimensions = [0] : (tensor<10xf32>, tensor<10xi32>, tensor, tensor) -> (tensor, tensor) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %217 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %218 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %219 = stablehlo.or %217, %218 : tensor + %220 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %221 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %222 = stablehlo.and %220, %221 : tensor + %223 = stablehlo.or %219, %222 : tensor + %224 = stablehlo.select %219, %amv_l, %amv_r : tensor, tensor + %225 = stablehlo.select %223, %ami_l, %ami_r : tensor, tensor + stablehlo.return %224, %225 : tensor, tensor + } + return %216#1, %154, %156 : tensor, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index 58c6f15f4..0335160f5 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -12,6 +12,13 @@ //! sliding-window schedule (window + pattern period), and a per-layer NoPE mask //! (SmolLM3). Llama / Qwen2 / Gemma2 keep their exact previous flag combinations, //! so their emitted graphs are byte-for-byte unchanged. +//! +//! A second dense pack (issue #499) rides the same flags with a config / naming +//! delta and no new emit: Seed-OSS / MiMo (the Qwen2 bias forward), InternLM3 +//! (plain / in-context `dynamic` RoPE), and ExaOne 3.x (llama3 RoPE with GPT-2-style +//! tensor names, selected by [`WeightScheme`]). Out-of-scope deltas an emit would +//! get wrong (interleaved RoPE, an o_proj / MLP bias, a non-SwiGLU activation, yarn +//! RoPE) are rejected rather than mis-emitted. /// How the RoPE inverse-frequency table is computed. Both kinds share the /// `outer(pos, inv_freq)` table build (see [`rope`](super::rope)); they differ @@ -64,6 +71,27 @@ pub struct QkNorm { pub one_plus: bool, } +/// The checkpoint tensor-naming scheme (issue #499). Almost every Llama-family +/// checkpoint uses the standard HF layout +/// (`model.layers.{i}.self_attn.q_proj.weight`, `model.embed_tokens.weight`, +/// `model.norm.weight`); ExaOne 3.x instead keeps the original GPT-2-style names +/// (`transformer.h.{i}.attn.attention.q_proj.weight`, `transformer.wte.weight`, +/// `transformer.ln_f.weight`, and a `c_fc_0` / `c_fc_1` / `c_proj` gated MLP). The +/// scheme is a loader-only concern: it maps the emitter's fixed arg order to the +/// checkpoint's tensor names (`weight_names` in [`weight_names`](crate::weight_names)), +/// so it never changes an emitted graph and two configs that differ only in scheme +/// emit byte-for-byte identical StableHLO. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum WeightScheme { + /// Standard HF Llama-family names (Llama, Qwen2, Gemma2, ERNIE-4.5, Seed-OSS, + /// MiMo, InternLM3, ...). + #[default] + Llama, + /// ExaOne 3.x GPT-2-style names (`transformer.h.{i}...`, gated MLP `c_fc_0` / + /// `c_fc_1` / `c_proj`, `out_proj` attention output). + Exaone, +} + /// MLX affine weight quantization (`config.json` `quantization`). The linear / /// embedding `*.weight` tensors are stored packed as `U32` with companion /// `*.scales` / `*.biases`; the loader dequantizes them to f32 as @@ -138,6 +166,12 @@ pub struct Config { /// Per-layer NoPE mask (SmolLM3): `use_rope_layers[li] == false` skips RoPE on /// that layer (`no_rope_layers`). `None` applies RoPE on every layer. pub use_rope_layers: Option>, + /// Checkpoint tensor-naming scheme (issue #499). Loader-only: it selects how + /// [`weight_names`](crate::weight_names) maps the emitter's arg order onto the + /// checkpoint tensors, so it never affects the emitted graph. `Llama` (the + /// default) is the standard HF layout; `Exaone` is ExaOne 3.x's GPT-2-style + /// names. + pub weight_scheme: WeightScheme, } impl Config { @@ -174,6 +208,7 @@ impl Config { sliding_window: None, sliding_pattern: 2, use_rope_layers: None, + weight_scheme: WeightScheme::Llama, } } @@ -191,6 +226,31 @@ impl Config { serde_json::from_str(s).map_err(|e| format!("parse config.json: {e}"))?; let model_type = v.get("model_type").and_then(serde_json::Value::as_str); + let bool_field = |k: &str| v.get(k).and_then(serde_json::Value::as_bool); + + // ExaOne 3.x keeps GPT-2-style tensor names; every other supported family + // uses the standard HF Llama layout. Loader-only (see [`WeightScheme`]), so + // it never changes the emitted graph. + let weight_scheme = if model_type == Some("exaone") { + WeightScheme::Exaone + } else { + WeightScheme::Llama + }; + + // Interleaved (GPT-J-style) RoPE is a distinct emitter delta from the half- + // split RoPE the supported families use: the pairs rotated are (2i, 2i+1), + // not (i, i+d/2). ERNIE-4.5 (`rotate_half` over `x[..., 0::2]` / + // `x[..., 1::2]`) is such an arch, so despite looking like a plain-RoPE Llama + // in config.json it is rejected here rather than mis-emitted (a half-split + // emit is close but wrong) pending an interleaved-RoPE variant. + if model_type == Some("ernie4_5") { + return Err( + "the OpenXLA emitter uses half-split RoPE; ERNIE-4.5 (model_type \ + ernie4_5) uses interleaved (GPT-J-style) RoPE, which is a follow-up \ + (an interleaved-RoPE emit variant or a load-time q/k permutation)" + .to_string(), + ); + } // Tied (share `embed` for the head) vs untied (separate `lm_head.weight`). // HF `PretrainedConfig` defaults this to `true`, so an absent field means @@ -235,12 +295,26 @@ impl Config { .and_then(serde_json::Value::as_u64) .map(|x| x as usize) }; + // Some arches use alternate field names (ExaOne 3.x: `num_layers`, + // `layer_norm_epsilon`), so these read the first present of a key list. + let u_any = |keys: &[&str]| -> Result { + keys.iter() + .find_map(|k| v.get(*k).and_then(serde_json::Value::as_u64)) + .map(|x| x as usize) + .ok_or_else(|| format!("config.json missing integer among {keys:?}")) + }; + let f_any = |keys: &[&str]| -> Result { + keys.iter() + .find_map(|k| v.get(*k).and_then(serde_json::Value::as_f64)) + .ok_or_else(|| format!("config.json missing number among {keys:?}")) + }; let hidden = u("hidden_size")?; let n_q = u("num_attention_heads")?; // head_dim is explicit in recent configs; otherwise it is hidden / heads. let head_dim = ou("head_dim").unwrap_or(hidden / n_q.max(1)); - let n_layers = u("num_hidden_layers")?; + // ExaOne 3.x uses `num_layers` in place of `num_hidden_layers`. + let n_layers = u_any(&["num_hidden_layers", "num_layers"])?; // Architecture-family flags, defaulted to the Llama baseline and overridden // per model_type below. @@ -271,10 +345,11 @@ impl Config { Some("llama") => { // A `llama` checkpoint with attention bias would need the Qwen2 bias // emit, untested here; reject rather than emit an unvalidated graph. - if v.get("attention_bias").and_then(serde_json::Value::as_bool) == Some(true) { + if bool_field("attention_bias") == Some(true) { return Err( "the OpenXLA emitter does not support a `llama` checkpoint with \ - attention_bias = true (only Qwen2 carries a q/k/v bias here)" + attention_bias = true (only the bias-bearing dense arches carry a \ + q/k/v bias here)" .to_string(), ); } @@ -375,16 +450,69 @@ impl Config { } rope_local_base = of64("rope_local_base_freq"); } + // issue #499 dense pack: each maps to a proven Llama / Qwen2 forward with + // a config / naming delta and no new emit. + Some("seed_oss") | Some("mimo") | Some("internlm3") => { + // Bias-bearing dense forwards: Seed-OSS / MiMo expose the q/k/v bias + // as `attention_bias`, InternLM3 as `qkv_bias`. MiMo's config + // `sliding_window` is deliberately not read (served globally, as for + // Qwen2); the rope block serves Seed-OSS `default` / InternLM3 + // `dynamic` as plain RoPE. + qkv_bias = bool_field("attention_bias") == Some(true) + || bool_field("qkv_bias") == Some(true); + } + Some("exaone") => { + // ExaOne 3.x: llama3-RoPE Llama with GPT-2-style tensor names (see + // `weight_scheme`) and the `num_layers` / `layer_norm_epsilon` + // alternate config field names read elsewhere. + } other => { return Err(format!( - "the OpenXLA emitter supports the Llama, Qwen2, Qwen3, Gemma1/2/3, \ - SmolLM3, and OLMo2/3 architectures; config.json model_type = {other:?}" + "the OpenXLA emitter supports the dense architectures Llama, Qwen2, \ + Qwen3, Gemma1/2/3, SmolLM3, OLMo2/3, Seed-OSS, MiMo, InternLM3, and \ + ExaOne; config.json model_type = {other:?} (MoE / MLA / fused-QKV / \ + interleaved-RoPE / novel-activation variants are follow-ups)" + )); + } + } + + // Out-of-scope deltas an emit would get wrong are rejected rather than + // silently dropped (issue #499): the o_proj bias and the MLP bias have no + // emit, and a non-SwiGLU activation on a non-GeGLU family would be + // mis-emitted. Gemma's GeGLU is driven by its own `mlp_geglu` flag, so the + // activation guard skips it. + if bool_field("attention_out_bias") == Some(true) { + return Err( + "the OpenXLA emitter has no attention output (o_proj) bias; \ + config.json attention_out_bias = true is a follow-up" + .to_string(), + ); + } + if bool_field("mlp_bias") == Some(true) { + return Err( + "the OpenXLA emitter has no MLP bias; config.json mlp_bias = true is a follow-up" + .to_string(), + ); + } + if !mlp_geglu { + let act = v + .get("hidden_act") + .or_else(|| v.get("activation_function")) + .and_then(serde_json::Value::as_str); + if let Some(a) = act + && a != "silu" + { + return Err(format!( + "the OpenXLA emitter emits a SwiGLU (silu) MLP for this architecture; \ + config.json activation = {a:?} is unsupported" )); } } - // rope_scaling is optional: absent -> plain RoPE; present -> only the llama3 - // scheme is supported (yarn, e.g. OLMo3 at full size, is a follow-up). + // rope_scaling is optional: absent -> plain RoPE (Qwen2.5, plain Llama). + // When present, the supported schemes are llama3 (scaled) and, served as + // plain RoPE, the `default` (identity) and in-context `dynamic` types; + // anything else (e.g. yarn, which OLMo3 uses at full size) is a follow-up. let rope = match v.get("rope_scaling") { None | Some(serde_json::Value::Null) => RopeScaling::Plain, Some(scaling) => { @@ -392,32 +520,45 @@ impl Config { .get("rope_type") .or_else(|| scaling.get("type")) .and_then(serde_json::Value::as_str); - if rope_type != Some("llama3") { - return Err(format!( - "the OpenXLA emitter supports plain RoPE and llama3 RoPE scaling; \ - config.json rope_scaling.rope_type = {rope_type:?} (e.g. yarn is a \ - follow-up)" - )); - } - let sf = |k: &str| -> Result { - scaling - .get(k) - .and_then(serde_json::Value::as_f64) - .ok_or_else(|| format!("config.json rope_scaling missing number `{k}`")) - }; - let orig_ctx = scaling - .get("original_max_position_embeddings") - .and_then(serde_json::Value::as_u64) - .map(|x| x as usize) - .ok_or_else(|| { - "config.json rope_scaling missing `original_max_position_embeddings`" - .to_string() - })?; - RopeScaling::Llama3 { - factor: sf("factor")?, - low_freq_factor: sf("low_freq_factor")?, - high_freq_factor: sf("high_freq_factor")?, - orig_ctx, + match rope_type { + // `default` is HF's identity rope (Seed-OSS). `dynamic` NTK + // (InternLM2/3) is identity within the original context and only + // rescales beyond it, so short / in-context generation is served + // as plain RoPE here (both use `rope_theta`); the long-context + // NTK rescale is a follow-up. + Some("default") | Some("dynamic") => RopeScaling::Plain, + Some("llama3") => { + let sf = |k: &str| -> Result { + scaling + .get(k) + .and_then(serde_json::Value::as_f64) + .ok_or_else(|| { + format!("config.json rope_scaling missing number `{k}`") + }) + }; + let orig_ctx = scaling + .get("original_max_position_embeddings") + .and_then(serde_json::Value::as_u64) + .map(|x| x as usize) + .ok_or_else(|| { + "config.json rope_scaling missing \ + `original_max_position_embeddings`" + .to_string() + })?; + RopeScaling::Llama3 { + factor: sf("factor")?, + low_freq_factor: sf("low_freq_factor")?, + high_freq_factor: sf("high_freq_factor")?, + orig_ctx, + } + } + other => { + return Err(format!( + "the OpenXLA emitter supports plain / default / (in-context) dynamic \ + RoPE and llama3 RoPE scaling; config.json rope_scaling.rope_type = \ + {other:?} (e.g. yarn is a follow-up)" + )); + } } } }; @@ -429,7 +570,8 @@ impl Config { n_q, n_kv: u("num_key_value_heads")?, head_dim, - eps: f("rms_norm_eps")? as f32, + // ExaOne 3.x uses `layer_norm_epsilon` in place of `rms_norm_eps`. + eps: f_any(&["rms_norm_eps", "layer_norm_epsilon"])? as f32, rope_theta: f("rope_theta")?, vocab: u("vocab_size")?, rope, @@ -448,6 +590,7 @@ impl Config { sliding_window, sliding_pattern, use_rope_layers, + weight_scheme, }) } @@ -526,3 +669,144 @@ impl Config { matches!(self.norm_style, NormStyle::GemmaFf | NormStyle::OlmoPost) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// ERNIE-4.5 is rejected with a message naming its interleaved (GPT-J-style) + /// RoPE: it looks like a plain-RoPE Llama in config.json but its `rotate_half` + /// rotates the (2i, 2i+1) pairs, not the (i, i+d/2) halves the emitter uses, so + /// a half-split emit would be wrong. Deferred to an interleaved-RoPE follow-up. + #[test] + fn rejects_ernie4_5_interleaved_rope() { + let j = r#"{"model_type":"ernie4_5","hidden_size":1024,"intermediate_size":3072, + "num_hidden_layers":18,"num_attention_heads":16,"num_key_value_heads":2, + "head_dim":128,"rms_norm_eps":1e-5,"rope_theta":500000,"vocab_size":103424, + "tie_word_embeddings":true,"hidden_act":"silu","use_bias":false}"#; + let err = Config::from_json_str(j).expect_err("ernie4_5 is deferred"); + assert!( + err.contains("interleaved"), + "names the interleaved-RoPE reason: {err}" + ); + } + + /// Seed-OSS parses to a Qwen2-style bias forward: `attention_bias = true` turns + /// on the q/k/v bias, `rope_type = "default"` is served as plain RoPE, and it is + /// untied. `attention_out_bias = false` is accepted (only `true` is rejected). + #[test] + fn parses_seed_oss_as_qkv_bias_default_rope() { + let j = r#"{"model_type":"seed_oss","hidden_size":5120,"intermediate_size":27648, + "num_hidden_layers":64,"num_attention_heads":80,"num_key_value_heads":8, + "head_dim":128,"rms_norm_eps":1e-6,"rope_theta":1e7,"vocab_size":155136, + "tie_word_embeddings":false,"attention_bias":true,"attention_out_bias":false, + "rope_scaling":{"rope_type":"default"},"hidden_act":"silu"}"#; + let c = Config::from_json_str(j).expect("seed_oss parses"); + assert!(c.qkv_bias, "attention_bias=true -> q/k/v bias"); + assert_eq!(c.rope, RopeScaling::Plain, "rope_type default -> plain"); + assert!(!c.tie_word_embeddings, "seed_oss is untied"); + } + + /// MiMo parses to a Qwen2-style bias forward, and its config `sliding_window` + /// is ignored (served globally, as for Qwen2), so it parses to `None`. + #[test] + fn parses_mimo_qkv_bias_ignores_sliding_window() { + let j = r#"{"model_type":"mimo","hidden_size":4096,"intermediate_size":11008, + "num_hidden_layers":36,"num_attention_heads":32,"num_key_value_heads":8, + "head_dim":128,"rms_norm_eps":1e-5,"rope_theta":640000,"vocab_size":151680, + "tie_word_embeddings":false,"attention_bias":true,"sliding_window":32768, + "use_sliding_window":true,"hidden_act":"silu"}"#; + let c = Config::from_json_str(j).expect("mimo parses"); + assert!(c.qkv_bias); + assert_eq!(c.rope, RopeScaling::Plain); + assert_eq!( + c.sliding_window, None, + "non-gemma2 sliding_window is ignored" + ); + } + + /// InternLM3 parses to a plain-RoPE untied Llama: `rope_type = "dynamic"` is + /// served as plain (in-context), and `qkv_bias` drives the bias (false here). + #[test] + fn parses_internlm3_dynamic_rope_as_plain() { + let j = r#"{"model_type":"internlm3","hidden_size":4096,"intermediate_size":10240, + "num_hidden_layers":48,"num_attention_heads":32,"num_key_value_heads":2, + "head_dim":128,"rms_norm_eps":1e-5,"rope_theta":50000000,"vocab_size":128512, + "tie_word_embeddings":false,"qkv_bias":false, + "rope_scaling":{"rope_type":"dynamic","factor":6.0},"hidden_act":"silu"}"#; + let c = Config::from_json_str(j).expect("internlm3 parses"); + assert_eq!(c.rope, RopeScaling::Plain, "dynamic -> plain (in-context)"); + assert!(!c.qkv_bias, "qkv_bias=false"); + assert!(!c.tie_word_embeddings); + // A `qkv_bias = true` internlm3 turns the bias on. + let biased = j.replace("\"qkv_bias\":false", "\"qkv_bias\":true"); + assert!(Config::from_json_str(&biased).unwrap().qkv_bias); + } + + /// ExaOne 3.x parses to a llama3-RoPE tied Llama with the ExaOne weight scheme, + /// reading the alternate field names (`num_layers`, `layer_norm_epsilon`). + #[test] + fn parses_exaone_alt_fields_and_scheme() { + let j = r#"{"model_type":"exaone","hidden_size":2560,"intermediate_size":7168, + "num_layers":30,"num_attention_heads":32,"num_key_value_heads":8,"head_dim":80, + "layer_norm_epsilon":1e-5,"rope_theta":1000000,"vocab_size":102400, + "tie_word_embeddings":true,"activation_function":"silu", + "rope_scaling":{"rope_type":"llama3","factor":8.0,"low_freq_factor":1.0, + "high_freq_factor":4.0,"original_max_position_embeddings":8192}}"#; + let c = Config::from_json_str(j).expect("exaone parses"); + assert_eq!(c.weight_scheme, WeightScheme::Exaone); + assert_eq!(c.n_layers, 30, "num_layers -> n_layers"); + assert_eq!(c.eps, 1e-5, "layer_norm_epsilon -> eps"); + assert_eq!(c.head_dim, 80); + assert!(c.tie_word_embeddings); + assert!(matches!(c.rope, RopeScaling::Llama3 { factor, .. } if factor == 8.0)); + } + + /// Unsupported deltas are rejected with a clear message rather than mis-emitted: + /// an attention output bias, an MLP bias, a non-SwiGLU activation, an + /// unsupported rope type (yarn), and an unsupported `model_type`. + #[test] + fn rejects_out_of_scope_dense_deltas() { + let base = |extra: &str| { + format!( + r#"{{"model_type":"seed_oss","hidden_size":8,"intermediate_size":16, + "num_hidden_layers":2,"num_attention_heads":2,"num_key_value_heads":1, + "rms_norm_eps":1e-6,"rope_theta":1e4,"vocab_size":10{extra}}}"# + ) + }; + assert!( + Config::from_json_str(&base(",\"attention_out_bias\":true")) + .unwrap_err() + .contains("o_proj"), + "o_proj bias rejected" + ); + assert!( + Config::from_json_str(&base(",\"mlp_bias\":true")) + .unwrap_err() + .contains("MLP bias"), + "mlp bias rejected" + ); + assert!( + Config::from_json_str(&base(",\"hidden_act\":\"gelu\"")) + .unwrap_err() + .contains("SwiGLU"), + "non-silu activation rejected" + ); + assert!( + Config::from_json_str(&base(",\"rope_scaling\":{\"rope_type\":\"yarn\"}")) + .unwrap_err() + .contains("yarn"), + "yarn rope rejected" + ); + // An architecture the emitter cannot reproduce (MoE / MLA glm4 variant). + let glm = r#"{"model_type":"glm4_moe_lite","hidden_size":8,"intermediate_size":16, + "num_hidden_layers":2,"num_attention_heads":2,"num_key_value_heads":1, + "rms_norm_eps":1e-5,"rope_theta":1e4,"vocab_size":10}"#; + assert!( + Config::from_json_str(glm) + .unwrap_err() + .contains("model_type"), + "unsupported model_type rejected" + ); + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 698a577b4..61fb51843 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -43,7 +43,7 @@ mod config; mod model; mod rope; -pub(crate) use config::Config; +pub(crate) use config::{Config, NormStyle, QkNorm, WeightScheme}; // `resolve_precision` and the precision-taking `*_with` emit variants are consumed // by the IREE execution path; the f32-default `emit_*` wrappers by the byte-exact // regression tests. Which set is live depends on the build cfg (`iree` vs `test`), @@ -107,6 +107,7 @@ mod tests { sliding_window: None, sliding_pattern: 2, use_rope_layers: None, + weight_scheme: super::WeightScheme::Llama, } } @@ -145,6 +146,7 @@ mod tests { sliding_window, sliding_pattern: 2, use_rope_layers: None, + weight_scheme: super::WeightScheme::Llama, } } @@ -642,6 +644,28 @@ mod tests { std::fs::write(&out_path, emit_prefill(&cfg, false)).expect("write MLXCEL_DUMP_OUT"); } + /// Opt-in graph dump for the issue #499 dense-pack execution check. Ignored by + /// default; with `MLXCEL_DUMP_CONFIG` (a checkpoint `config.json`) and + /// `MLXCEL_DUMP_DIR` set, it writes the host-sampled prefill + decode logits + /// graphs (`prefill_logits.mlir` / `decode_logits.mlir`) from the REAL config, + /// so `spike/openxla/arch_execution_check.py` can compile them with IREE and + /// drive a token-exact continuation against an HF fp32 oracle. Pure Rust / + /// scoped, so the check never needs the `iree` cargo feature. + #[test] + #[ignore = "opt-in: writes prefill/decode graphs to disk for the spike/openxla execution check"] + fn dump_dense_pack_graphs_for_execution_check() { + let cfg_path = + std::env::var("MLXCEL_DUMP_CONFIG").expect("set MLXCEL_DUMP_CONFIG to a config.json"); + let dir = std::env::var("MLXCEL_DUMP_DIR").expect("set MLXCEL_DUMP_DIR to an output dir"); + let text = std::fs::read_to_string(&cfg_path).expect("read MLXCEL_DUMP_CONFIG"); + let cfg = Config::from_json_str(&text).expect("parse config.json"); + let dir = std::path::Path::new(&dir); + std::fs::write(dir.join("prefill_logits.mlir"), emit_prefill(&cfg, false)) + .expect("write prefill_logits.mlir"); + std::fs::write(dir.join("decode_logits.mlir"), emit_decode(&cfg, false)) + .expect("write decode_logits.mlir"); + } + /// Plain RoPE base frequencies are the textbook `1 / theta^(2i/head_dim)` /// (Qwen2), distinct from the llama3-scaled table. #[test] @@ -725,6 +749,7 @@ mod tests { sliding_window: None, sliding_pattern: 2, use_rope_layers: None, + weight_scheme: super::WeightScheme::Llama, } } diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index 6ca634d80..44c9500a6 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -55,6 +55,11 @@ use safetensors::{Dtype, SafeTensors}; use crate::emitter::{ Config, emit_decode_ragged_with, emit_decode_with, emit_prefill_with, resolve_precision, }; +// Checkpoint tensor names in the emitter's exact arg order, per the model config's +// naming scheme (issue #499). Lives in a pure-Rust module so the ordering and each +// architecture's names are unit-tested without the `iree` feature; the loader +// below consumes the result unchanged. +use crate::weight_names::weight_names; use crate::weights::{bf16_to_f32, dequantize_affine, f16_to_f32, f32_le_to_f32}; /// Prefill bucket baked into the emitted `prefill` graph (`tensor<256xi32>`, @@ -124,69 +129,6 @@ unsafe extern "C" { /// load (any `b_max` is emittable; the worker selects from this set). pub(crate) const RAGGED_B_VALUES: &[usize] = &[4, 8]; -/// The weight names in the emitter's exact arg order: embed, final_norm, then — -/// for an untied checkpoint (`tie_word_embeddings = false`) — `lm_head.weight`, -/// then per layer the base weights and the arch-conditional extras. The base order -/// is down, gate, input_layernorm, post_attention_layernorm, up, wk, wo, wq, wv; -/// then, matching `take_layer_weights` in `emitter/model.rs` exactly, the k/q/v -/// biases (`qkv_bias`, Qwen2), the q/k norms (`qk_norm`, Qwen3 / Gemma3 / OLMo2/3), -/// and the feed-forward norms. `input_layernorm` is skipped for the OLMo reordered -/// post-norm (which has none). Every knob comes from the model config so the buffer -/// order lines up with the emitted graph args. -fn weight_names(cfg: &Config) -> Vec { - let mut names = vec![ - "model.embed_tokens.weight".to_string(), - "model.norm.weight".to_string(), - ]; - // Untied LM head: a separate `lm_head.weight` follows `final_norm`, matching - // the `params['lm_head']` arg the emitter takes in the same position. - if !cfg.tie_word_embeddings { - names.push("lm_head.weight".to_string()); - } - for i in 0..cfg.n_layers { - let p = format!("model.layers.{i}."); - names.push(format!("{p}mlp.down_proj.weight")); - names.push(format!("{p}mlp.gate_proj.weight")); - // input_layernorm: present unless the reordered (OLMo) post-norm drops it. - if cfg.has_input_norm() { - names.push(format!("{p}input_layernorm.weight")); - } - for suf in [ - "post_attention_layernorm.weight", - "mlp.up_proj.weight", - "self_attn.k_proj.weight", - "self_attn.o_proj.weight", - "self_attn.q_proj.weight", - "self_attn.v_proj.weight", - ] { - names.push(format!("{p}{suf}")); - } - // Qwen2 q/k/v projection biases, in the same k/q/v order the emitter adds. - if cfg.qkv_bias { - for suf in [ - "self_attn.k_proj.bias", - "self_attn.q_proj.bias", - "self_attn.v_proj.bias", - ] { - names.push(format!("{p}{suf}")); - } - } - // q/k norms (Qwen3 / Gemma3 per-head, OLMo2/3 flat), q then k. - if cfg.qk_norm.is_some() { - names.push(format!("{p}self_attn.q_norm.weight")); - names.push(format!("{p}self_attn.k_norm.weight")); - } - // Feed-forward norms: Gemma2/3 add pre AND post; OLMo2/3 add post only. - if cfg.has_pre_ff_norm() { - names.push(format!("{p}pre_feedforward_layernorm.weight")); - } - if cfg.has_post_ff_norm() { - names.push(format!("{p}post_feedforward_layernorm.weight")); - } - } - names -} - /// Locate the IREE distribution: a runtime `IREE_DIST` override first, else the /// path baked at build time (the dist whose runtime is linked into this binary). /// Only the dist (CPU/vulkan) build uses this; the cuda build links a source-built diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index c2ec15930..9a4a43c1a 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -73,6 +73,13 @@ mod sampler; #[cfg_attr(not(feature = "iree"), allow(dead_code))] mod weights; +// Per-architecture checkpoint tensor naming (#449 M3 Stage 2d; generalized in +// #499). Pure Rust; present under `iree` (the loader orders weights with it) and +// under `test` (the naming schemes are unit-tested without the IREE runtime). +#[cfg(any(feature = "iree", test))] +#[cfg_attr(not(feature = "iree"), allow(dead_code))] +mod weight_names; + // Rust-native StableHLO emitter (#449 M3 Stage 2d, ported from the #451 spike). // Pure Rust; present under `iree` (the engine emits its graphs from config.json at // load) and under `test` (the byte-exact regression runs without the IREE diff --git a/src/lib/mlxcel-xla/src/validation.rs b/src/lib/mlxcel-xla/src/validation.rs index ba5d2d50f..c7fe3978e 100644 --- a/src/lib/mlxcel-xla/src/validation.rs +++ b/src/lib/mlxcel-xla/src/validation.rs @@ -357,9 +357,51 @@ pub(crate) static LLAMA_3_2_1B: ArchFixture = ArchFixture { ], }; +/// One dense-pack family (issue #499). Every listed architecture reuses the +/// already-proven Llama or Qwen2 forward up to a config / naming delta, so its +/// emitted graph is a proven graph and the frozen goldens are trusted by +/// construction (the emit is exercised absolutely by [`LLAMA_3_2_1B`] and +/// relatively by `dense_pack_families_reuse_proven_graphs`). The fixtures use +/// small synthetic dims (real switches, tiny hidden / layers) so the goldens stay +/// small; the real checkpoint `config.json` parsing is covered by the +/// `config::tests` parse assertions. +macro_rules! dense_fixture { + ($ident:ident, $arch:literal) => { + pub(crate) static $ident: ArchFixture = ArchFixture { + arch: $arch, + config_json: include_str!(concat!("../assets/", $arch, "/config.json")), + graphs: &[ + GraphFixture { + kind: GraphKind::Decode { sample: true }, + golden_name: "decode.mlir", + golden: include_str!(concat!("../assets/", $arch, "/decode.mlir")), + }, + GraphFixture { + kind: GraphKind::Prefill { sample: true }, + golden_name: "prefill.mlir", + golden: include_str!(concat!("../assets/", $arch, "/prefill.mlir")), + }, + ], + }; + }; +} + +// Seed-OSS: q/k/v projection bias (from `attention_bias`), untied, `default` rope +// type served as plain — the proven Qwen2 bias forward with standard names. +dense_fixture!(SEED_OSS, "seed_oss"); +// MiMo: q/k/v projection bias, untied, plain RoPE; its config `sliding_window` is +// served globally (as for Qwen2), so it parses to `sliding_window = None`. +dense_fixture!(MIMO, "mimo"); +// InternLM3: standard names, untied, `dynamic` rope served as plain (in-context). +dense_fixture!(INTERNLM3, "internlm3"); +// ExaOne 3.x: llama3 RoPE, tied, GPT-2-style names (the `Exaone` weight scheme) +// and the `num_layers` / `layer_norm_epsilon` alternate config fields. +dense_fixture!(EXAONE, "exaone"); + /// Every registered structural fixture. Append a family here to add it to the /// byte-exact gate; see the module docs for the freeze workflow. -pub(crate) static REGISTERED: &[&ArchFixture] = &[&LLAMA_3_2_1B]; +pub(crate) static REGISTERED: &[&ArchFixture] = + &[&LLAMA_3_2_1B, &SEED_OSS, &MIMO, &INTERNLM3, &EXAONE]; /// A golden-less structural fixture (issue #497): a small synthetic `config.json` /// for a dense family plus the signature every emitted shared-core graph must @@ -480,6 +522,83 @@ mod tests { } } + /// The dense pack (issue #499) reuses the proven forward: each family's config + /// emits StableHLO byte-for-byte identical to a proven `llama` / `qwen2` + /// reference carrying the SAME switches and dimensions, across the single, + /// prefill, and ragged graph kinds. This is the correctness anchor for the + /// families without a real-checkpoint execution run in-agent: their emit is + /// literally the Llama / Qwen2 emit the backend already runs token-exact, so + /// the delta is confined to config parsing (asserted in `config::tests`) and + /// tensor naming (asserted in `weight_names::tests`). The references match each + /// fixture's small synthetic dims / eps / theta exactly (rope tables and the + /// eps constant must line up for a byte match). + #[test] + fn dense_pack_families_reuse_proven_graphs() { + // (fixture, an equivalent config expressed via a proven model_type). + let cases: &[(&ArchFixture, &str)] = &[ + // Seed-OSS == untied Qwen2 (q/k/v bias, plain RoPE). + ( + &SEED_OSS, + r#"{"model_type":"qwen2","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":2,"rms_norm_eps":1e-6,"rope_theta":10000000.0, + "vocab_size":10,"tie_word_embeddings":false}"#, + ), + // MiMo == untied Qwen2 (q/k/v bias, plain RoPE), sliding_window ignored. + ( + &MIMO, + r#"{"model_type":"qwen2","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":2,"rms_norm_eps":1e-5,"rope_theta":640000.0, + "vocab_size":10,"tie_word_embeddings":false}"#, + ), + // InternLM3 == untied plain-RoPE Llama (dynamic served as plain). + ( + &INTERNLM3, + r#"{"model_type":"llama","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":2,"rms_norm_eps":1e-5,"rope_theta":50000000.0, + "vocab_size":10,"tie_word_embeddings":false}"#, + ), + // ExaOne 3.x == tied llama3-RoPE Llama (the weight scheme is loader-only). + ( + &EXAONE, + r#"{"model_type":"llama","hidden_size":8,"num_attention_heads":2, + "num_key_value_heads":1,"head_dim":4,"intermediate_size":16, + "num_hidden_layers":2,"rms_norm_eps":1e-5,"rope_theta":1000000.0, + "vocab_size":10,"tie_word_embeddings":true,"rope_scaling":{"rope_type":"llama3", + "factor":8.0,"low_freq_factor":1.0,"high_freq_factor":4.0, + "original_max_position_embeddings":8192}}"#, + ), + ]; + for (fx, ref_json) in cases { + let fam = Config::from_json_str(fx.config_json) + .unwrap_or_else(|e| panic!("{}: {e}", fx.arch)); + let refc = Config::from_json_str(ref_json) + .unwrap_or_else(|e| panic!("{} reference: {e}", fx.arch)); + let pairs = [ + ("decode", emit_decode(&fam, true), emit_decode(&refc, true)), + ( + "prefill", + emit_prefill(&fam, false), + emit_prefill(&refc, false), + ), + ( + "ragged", + emit_decode_ragged(&fam, 4, false), + emit_decode_ragged(&refc, 4, false), + ), + ]; + for (name, got, want) in pairs { + assert_eq!( + got, want, + "{}: emitted {name} is not byte-identical to its proven equivalent", + fx.arch + ); + } + } + } + /// The guard actually catches drift: a wrong golden fails and the report /// localizes the first differing line, so a downstream emitter change that /// shifts a graph is caught with a pointer, not a wall of text. diff --git a/src/lib/mlxcel-xla/src/weight_names.rs b/src/lib/mlxcel-xla/src/weight_names.rs new file mode 100644 index 000000000..99f1b8ef0 --- /dev/null +++ b/src/lib/mlxcel-xla/src/weight_names.rs @@ -0,0 +1,378 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Checkpoint tensor names in the emitter's weight arg order (issue #449 M3 Stage +//! 2d; generalized to per-architecture naming schemes in #499, and to the qk-norm / +//! Gemma dense pack in #497). +//! +//! The IREE loader ([`iree`](crate::iree)) reads each weight the emitted graph +//! takes as an argument, in the emitter's exact arg order, from the checkpoint's +//! safetensors. [`weight_names`] produces that ordered name list from the model +//! [`Config`], mirroring `take_lm_head` / `take_layer_weights` in `emitter/model.rs` +//! exactly: embed, final_norm, then — for an untied checkpoint +//! (`tie_word_embeddings = false`) — the LM head, then per layer down, gate, +//! `input_layernorm` (unless the OLMo reordered post-norm drops it), +//! post_attention_layernorm, up, wk, wo, wq, wv, then the arch-conditional extras in +//! the emitter's order: the k/q/v projection biases (`qkv_bias`), the q/k norms +//! (`qk_norm`), and the feed-forward norms (`pre_feedforward_layernorm` for Gemma +//! 2/3, `post_feedforward_layernorm` for Gemma 2/3 and OLMo 2/3). Llama / Qwen2 / +//! Gemma2 have none of the new extras, so their name lists are byte-for-byte +//! unchanged. +//! +//! Only the *names* vary by [`WeightScheme`] (issue #499): almost every checkpoint +//! uses the standard HF Llama layout, while ExaOne 3.x keeps GPT-2-style names. +//! The scheme never changes the emitted graph, so a family that differs only in +//! naming reuses the proven Llama / Qwen2 forward and its structural goldens +//! unchanged; this module is the one place the naming delta lives. +//! +//! Pure Rust (no IREE), so the ordering is unit-tested without the `iree` feature +//! (the loader that consumes it is `iree`-gated). + +use crate::emitter::{Config, WeightScheme}; + +/// The tensor-name pieces for one [`WeightScheme`]: the top-level embed / +/// final-norm / (untied) LM-head names, the per-layer prefix stem +/// (`"{stem}{i}."`), and each per-layer suffix, emitted in the emitter's arg order +/// by [`weight_names`] under the same config gates as `take_layer_weights`. +struct SchemeNames { + /// Token-embedding weight name. + embed: &'static str, + /// Final RMSNorm weight name. + final_norm: &'static str, + /// Untied LM-head weight name (used only when `tie_word_embeddings = false`). + lm_head: &'static str, + /// Per-layer prefix stem; the full prefix is `format!("{stem}{i}.")`. + layer_stem: &'static str, + /// MLP down projection. + down: &'static str, + /// MLP gate projection (the activation input). + gate: &'static str, + /// `input_layernorm` (skipped for the OLMo reordered post-norm). + input_layernorm: &'static str, + /// `post_attention_layernorm`. + post_attention_layernorm: &'static str, + /// MLP up projection. + up: &'static str, + /// Attention key projection. + k_proj: &'static str, + /// Attention output projection. + o_proj: &'static str, + /// Attention query projection. + q_proj: &'static str, + /// Attention value projection. + v_proj: &'static str, + /// Attention key projection bias (used only for a `qkv_bias` arch). + k_bias: &'static str, + /// Attention query projection bias (used only for a `qkv_bias` arch). + q_bias: &'static str, + /// Attention value projection bias (used only for a `qkv_bias` arch). + v_bias: &'static str, + /// Query RMSNorm weight (used only for a `qk_norm` arch). + q_norm: &'static str, + /// Key RMSNorm weight (used only for a `qk_norm` arch). + k_norm: &'static str, + /// Pre-feed-forward norm (used only for the Gemma four-norm layer). + pre_ff_norm: &'static str, + /// Post-feed-forward norm (used only for Gemma 2/3 and OLMo 2/3). + post_ff_norm: &'static str, +} + +/// The name pieces for `scheme`. `Llama` reproduces the standard HF layout the +/// loader has always used (byte-for-byte, so Llama / Qwen2 / Gemma2 loading is +/// unchanged); `Exaone` is ExaOne 3.x's GPT-2-style layout, verified against the +/// checkpoint's `modeling_exaone.py` (gated MLP `c_proj(act(c_fc_0(x)) * +/// c_fc_1(x))`, so `c_fc_0` is the gate and `c_fc_1` the up projection; attention +/// under `attn.attention.*` with `out_proj` as o_proj). +fn scheme_names(scheme: WeightScheme) -> SchemeNames { + match scheme { + WeightScheme::Llama => SchemeNames { + embed: "model.embed_tokens.weight", + final_norm: "model.norm.weight", + lm_head: "lm_head.weight", + layer_stem: "model.layers.", + down: "mlp.down_proj.weight", + gate: "mlp.gate_proj.weight", + input_layernorm: "input_layernorm.weight", + post_attention_layernorm: "post_attention_layernorm.weight", + up: "mlp.up_proj.weight", + k_proj: "self_attn.k_proj.weight", + o_proj: "self_attn.o_proj.weight", + q_proj: "self_attn.q_proj.weight", + v_proj: "self_attn.v_proj.weight", + k_bias: "self_attn.k_proj.bias", + q_bias: "self_attn.q_proj.bias", + v_bias: "self_attn.v_proj.bias", + q_norm: "self_attn.q_norm.weight", + k_norm: "self_attn.k_norm.weight", + pre_ff_norm: "pre_feedforward_layernorm.weight", + post_ff_norm: "post_feedforward_layernorm.weight", + }, + WeightScheme::Exaone => SchemeNames { + embed: "transformer.wte.weight", + final_norm: "transformer.ln_f.weight", + lm_head: "lm_head.weight", + layer_stem: "transformer.h.", + down: "mlp.c_proj.weight", // down + gate: "mlp.c_fc_0.weight", // gate (activation input) + input_layernorm: "ln_1.weight", // in_ln + post_attention_layernorm: "ln_2.weight", // post_ln + up: "mlp.c_fc_1.weight", // up + k_proj: "attn.attention.k_proj.weight", + o_proj: "attn.attention.out_proj.weight", + q_proj: "attn.attention.q_proj.weight", + v_proj: "attn.attention.v_proj.weight", + // ExaOne 3.x carries none of the conditional extras (qkv_bias, qk_norm, + // and the feed-forward norms are all off for it), so the names below are + // never emitted; they keep the scheme mapping total and mirror the + // attention path where one exists. + k_bias: "attn.attention.k_proj.bias", + q_bias: "attn.attention.q_proj.bias", + v_bias: "attn.attention.v_proj.bias", + q_norm: "attn.attention.q_norm.weight", + k_norm: "attn.attention.k_norm.weight", + pre_ff_norm: "pre_feedforward_layernorm.weight", + post_ff_norm: "post_feedforward_layernorm.weight", + }, + } +} + +/// The checkpoint tensor names in the emitter's exact arg order for `cfg`, matching +/// `take_lm_head` / `take_layer_weights` in `emitter/model.rs`. The names follow +/// `cfg.weight_scheme`, and every per-layer knob (the untied head, the conditional +/// `input_layernorm`, the q/k/v biases, the q/k norms, the feed-forward norms) is +/// read from `cfg`, so Llama / Qwen2 / Gemma2 stay byte-for-byte unchanged. +pub(crate) fn weight_names(cfg: &Config) -> Vec { + let s = scheme_names(cfg.weight_scheme); + let mut names = vec![s.embed.to_string(), s.final_norm.to_string()]; + // Untied LM head: a separate head weight follows `final_norm`, matching the + // `params['lm_head']` arg the emitter takes in the same position. + if !cfg.tie_word_embeddings { + names.push(s.lm_head.to_string()); + } + for i in 0..cfg.n_layers { + let p = format!("{}{i}.", s.layer_stem); + let mut push = |suf: &str| names.push(format!("{p}{suf}")); + push(s.down); + push(s.gate); + // input_layernorm: present unless the reordered (OLMo) post-norm drops it. + if cfg.has_input_norm() { + push(s.input_layernorm); + } + push(s.post_attention_layernorm); + push(s.up); + push(s.k_proj); + push(s.o_proj); + push(s.q_proj); + push(s.v_proj); + // q/k/v projection biases (Qwen2 / Seed-OSS / MiMo / InternLM3), in the same + // k/q/v order `take_layer_weights` adds them to the emitted graph args. + if cfg.qkv_bias { + push(s.k_bias); + push(s.q_bias); + push(s.v_bias); + } + // q/k norms (Qwen3 / Gemma3 per-head, OLMo2/3 flat), q then k. + if cfg.qk_norm.is_some() { + push(s.q_norm); + push(s.k_norm); + } + // Feed-forward norms: Gemma2/3 add pre AND post; OLMo2/3 add post only. + if cfg.has_pre_ff_norm() { + push(s.pre_ff_norm); + } + if cfg.has_post_ff_norm() { + push(s.post_ff_norm); + } + } + names +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::emitter::{NormStyle, QkNorm}; + + /// A minimal config for the namer. Starts from the Llama-3.2-1B defaults (Plain + /// norm, no bias, no q/k norm) and overrides the fields `weight_names` reads: + /// `n_layers`, `tie_word_embeddings`, `qkv_bias`, `weight_scheme`. Tests that + /// exercise the norm placement / q-k norm set `norm_style` / `qk_norm` directly. + fn cfg(n_layers: usize, tie: bool, qkv_bias: bool, scheme: WeightScheme) -> Config { + let mut c = Config::llama_3_2_1b(); + c.n_layers = n_layers; + c.tie_word_embeddings = tie; + c.qkv_bias = qkv_bias; + c.weight_scheme = scheme; + c + } + + /// The Llama scheme reproduces the exact standard HF names the loader has + /// always used, so Llama / Qwen2 / Gemma2 loading is byte-for-byte unchanged. + #[test] + fn llama_scheme_matches_the_standard_hf_layout() { + let names = weight_names(&cfg(1, true, false, WeightScheme::Llama)); + assert_eq!( + names, + vec![ + "model.embed_tokens.weight", + "model.norm.weight", + "model.layers.0.mlp.down_proj.weight", + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.input_layernorm.weight", + "model.layers.0.post_attention_layernorm.weight", + "model.layers.0.mlp.up_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + ] + ); + } + + /// Untied adds exactly the LM head, right after `final_norm`, before the layers. + #[test] + fn untied_inserts_lm_head_after_final_norm() { + let names = weight_names(&cfg(1, false, false, WeightScheme::Llama)); + assert_eq!(names[2], "lm_head.weight"); + let fnorm = names.iter().position(|n| n == "model.norm.weight").unwrap(); + let lm = names.iter().position(|n| n == "lm_head.weight").unwrap(); + let l0 = names + .iter() + .position(|n| n.starts_with("model.layers.0.")) + .unwrap(); + assert!(fnorm < lm && lm < l0, "final_norm < lm_head < layer0"); + } + + /// A `qkv_bias` arch (Seed-OSS / MiMo / Qwen2) appends the three biases per + /// layer in k/q/v order, after the nine core tensors. + #[test] + fn qkv_bias_appends_three_biases_in_k_q_v_order() { + let names = weight_names(&cfg(1, false, true, WeightScheme::Llama)); + let idx = |n: &str| names.iter().position(|x| x == n).unwrap(); + let (v, bk, bq, bv) = ( + idx("model.layers.0.self_attn.v_proj.weight"), + idx("model.layers.0.self_attn.k_proj.bias"), + idx("model.layers.0.self_attn.q_proj.bias"), + idx("model.layers.0.self_attn.v_proj.bias"), + ); + assert!(v < bk && bk < bq && bq < bv, "wv < bk < bq < bv"); + } + + /// A `qk_norm` arch (Qwen3 / Gemma3 / OLMo2/3) appends q_norm then k_norm per + /// layer, after the projections and after any biases. + #[test] + fn qk_norm_appends_q_then_k_norm_after_projections() { + let mut c = cfg(1, true, false, WeightScheme::Llama); + c.qk_norm = Some(QkNorm { + per_head: true, + one_plus: false, + }); + let names = weight_names(&c); + let idx = |n: &str| names.iter().position(|x| x == n).unwrap(); + let (v, qn, kn) = ( + idx("model.layers.0.self_attn.v_proj.weight"), + idx("model.layers.0.self_attn.q_norm.weight"), + idx("model.layers.0.self_attn.k_norm.weight"), + ); + assert!(v < qn && qn < kn, "wv < q_norm < k_norm"); + } + + /// The Gemma four-norm layer keeps `input_layernorm` and adds the pre AND post + /// feed-forward norms (pre before post), matching `take_layer_weights`. + #[test] + fn gemma_ff_adds_pre_then_post_feedforward_norm() { + let mut c = cfg(1, true, false, WeightScheme::Llama); + c.norm_style = NormStyle::GemmaFf; + let names = weight_names(&c); + assert!( + names.iter().any(|n| n.ends_with("input_layernorm.weight")), + "Gemma keeps input_layernorm" + ); + let idx = |suf: &str| names.iter().position(|x| x.ends_with(suf)).unwrap(); + let pre = idx("pre_feedforward_layernorm.weight"); + let post = idx("post_feedforward_layernorm.weight"); + assert!( + pre < post, + "pre-feedforward norm before post-feedforward norm" + ); + } + + /// The OLMo reordered post-norm drops `input_layernorm` and adds only the + /// `post_feedforward_layernorm` (no pre), while keeping the q/k norm. + #[test] + fn olmo_post_norm_drops_input_norm_and_pre_ff() { + let mut c = cfg(1, false, false, WeightScheme::Llama); + c.norm_style = NormStyle::OlmoPost; + c.qk_norm = Some(QkNorm { + per_head: false, + one_plus: false, + }); + let names = weight_names(&c); + assert!( + !names.iter().any(|n| n.ends_with("input_layernorm.weight")), + "OLMo post-norm has no input_layernorm" + ); + assert!( + names + .iter() + .any(|n| n.ends_with("post_feedforward_layernorm.weight")), + "OLMo post-norm keeps the post-feedforward norm" + ); + assert!( + !names + .iter() + .any(|n| n.ends_with("pre_feedforward_layernorm.weight")), + "OLMo post-norm has no pre-feedforward norm" + ); + assert!( + names.iter().any(|n| n.ends_with("q_norm.weight")), + "OLMo keeps the flat q/k norm" + ); + } + + /// The ExaOne scheme maps the emitter's arg order onto the GPT-2-style names, + /// with the critical gated-MLP mapping (down←c_proj, gate←c_fc_0, up←c_fc_1) + /// and `out_proj` as the attention output; it is tied (no LM head). + #[test] + fn exaone_scheme_maps_gpt2_style_names() { + let names = weight_names(&cfg(1, true, false, WeightScheme::Exaone)); + assert_eq!( + names, + vec![ + "transformer.wte.weight", + "transformer.ln_f.weight", + "transformer.h.0.mlp.c_proj.weight", // down + "transformer.h.0.mlp.c_fc_0.weight", // gate + "transformer.h.0.ln_1.weight", // in_ln + "transformer.h.0.ln_2.weight", // post_ln + "transformer.h.0.mlp.c_fc_1.weight", // up + "transformer.h.0.attn.attention.k_proj.weight", // wk + "transformer.h.0.attn.attention.out_proj.weight", // wo + "transformer.h.0.attn.attention.q_proj.weight", // wq + "transformer.h.0.attn.attention.v_proj.weight", // wv + ] + ); + } + + /// The per-layer name count is stable: 9 core (+3 biases when qkv_bias), for + /// every scheme. + #[test] + fn per_layer_counts_scale_with_the_deltas() { + for scheme in [WeightScheme::Llama, WeightScheme::Exaone] { + let base = weight_names(&cfg(2, true, false, scheme)).len(); + assert_eq!(base, 2 + 2 * 9, "embed+norm + 2 layers * 9 core"); + let biased = weight_names(&cfg(2, true, true, scheme)).len(); + assert_eq!(biased - base, 2 * 3, "qkv_bias adds 3 per layer"); + } + } +}