from __future__ import annotations

import mlx.core as mx

from microscope.fuse2_mlx import Model, ModelArgs


def main():
    args = ModelArgs(
        model_type="fuse2",
        hidden_size=32,
        num_hidden_layers=2,
        intermediate_size=64,
        num_attention_heads=4,
        rms_norm_eps=1e-6,
        vocab_size=128,
        num_key_value_heads=2,
        max_position_embeddings=64,
        rope_theta=1000000.0,
        head_dim=8,
        tie_word_embeddings=True,
        experts_per_layer={"0": [0, 1]},
        expert_hidden_size=32,
        expert_intermediate_size=16,
    )
    model = Model(args)
    ids = mx.array([[1, 2, 3]])
    uncached = model(ids)
    cache = model.make_cache()
    model(ids[:, :-1], cache)
    cached = model(ids[:, -1:], cache)
    mx.eval(uncached, cached)
    error = float(mx.max(mx.abs(uncached[:, -1] - cached[:, -1])))
    if error > 1e-4:
        raise AssertionError(error)
    print({"status": "validated", "max_logit_error": error})


if __name__ == "__main__":
    main()