import mlx.core as mx

weights = mx.load("model-00001-of-00002.safetensors")
lm_head = weights["lm_head.weight"]
embed = weights["model.language_model.embed_tokens.weight"]

print("shapes:", lm_head.shape, embed.shape)
print("identical:", bool(mx.array_equal(lm_head, embed)))
print("max abs diff:", float(mx.abs(lm_head.astype(mx.float32) - embed.astype(mx.float32)).max()))