package model

import (
	"fmt"
	"math"
	"os"
	"reflect"
	"testing"

	"github.com/adalkiran/llama-nuts-and-bolts/src/common"
	"github.com/adalkiran/llama-nuts-and-bolts/src/dtype"
	"github.com/adalkiran/llama-nuts-and-bolts/src/ml"
)

/*
	This simulation was added to ensure all steps are done correctly
	(until only first attention layer).
*/

func testTransformer_Prepare(t *testing.T, skipCompareTestTensor bool, transformer *LlamaTransformer, inputTokens *ml.Tensor, startPos int) (actualInputTensor *ml.Tensor, actualFreqsCis *ml.Tensor, actualMask *ml.Tensor) {
	expectedInputTensorSize := []int{15, 4096}
	// Shortened form as corresponding indices [0, 1, 2, 4093, 4094, 4095]
	expectedInputTensorShortened := [][]float32{
		{2.6512e-04, -4.9973e-04, -5.8365e-04 /*...,*/, 3.8147e-03, 6.3419e-05, 1.1902e-03},
		{-1.6499e-04, -2.4319e-04, 1.6403e-04 /*...,*/, -1.5163e-04, 3.5095e-04, 7.3242e-04},
		{3.5095e-03, 7.2021e-03, 5.3406e-05 /*...,*/, -7.2479e-04, -1.0620e-02, 8.2779e-04},
		/*...,*/
		{-9.7656e-03, -3.4637e-03, 1.8616e-03 /*...,*/, -7.1411e-03, -4.3030e-03, 8.6060e-03},
		{-4.6158e-04, -3.9291e-04, -6.5863e-06 /*...,*/, -6.2561e-04, -5.0354e-04, 6.6757e-04},
		{-2.8687e-03, 3.8910e-03, -1.7357e-04 /*...,*/, 8.0872e-04, 5.0354e-04, 2.3041e-03},
	}

	expectedFreqsCisSize := []int{15, 64}

	expectedMaskSize := []int{15, 15}
	negInf := float32(math.Inf(-1))
	expectedMask := [][]float32{
		{0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, negInf, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, negInf, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, negInf, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, negInf, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, negInf},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
	}
	var err error
	if actualInputTensor, actualFreqsCis, actualMask, err = transformer.prepare(inputTokens, startPos); err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedInputTensorShortened, expectedInputTensorSize, actualInputTensor, common.THRESHOLD_F32, true); err != nil {
		t.Fatal(err)
	}

	if !skipCompareTestTensor && !reflect.DeepEqual(expectedFreqsCisSize, actualFreqsCis.Size) {
		t.Fatalf("expected size %v, but got %v", expectedFreqsCisSize, actualFreqsCis.Size)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedMask, expectedMaskSize, actualMask, common.THRESHOLD_F32, false); err != nil {
		t.Fatal(err)
	}
	return
}

func testTransformerBlock_AttnNorm_Forward(t *testing.T, skipCompareTestTensor bool, transformerBlock *LlamaTransformerBlock, x *ml.Tensor) *ml.Tensor {
	/*
		normalizedX, err := ltb.attn_norm.Forward(infContext, x)
	*/
	expectedAttnNormPartSize := []int{15, 4096}
	expectedAttnNormPart := [][]float32{
		{0.0328, -0.0617, -0.0721 /*...,*/, 0.4714, 0.0078, 0.1471},
		{-0.0515, -0.0759, 0.0512 /*...,*/, -0.0473, 0.1095, 0.2284},
		{0.3651, 0.7492, 0.0056 /*...,*/, -0.0754, -1.1047, 0.0861},
		/*...,*/
		{-0.8352, -0.2962, 0.1592 /*...,*/, -0.6107, -0.3680, 0.7360},
		{-0.1443, -0.1228, -0.0021 /*...,*/, -0.1956, -0.1574, 0.2087},
		{-0.3833, 0.5199, -0.0232 /*...,*/, 0.1081, 0.0673, 0.3079},
	}

	expectedAttnNormalizedXSize := []int{15, 4096}
	expectedAttnNormalizedX := [][]float32{
		{0.0015, -0.0116, -0.0301 /*...,*/, 0.0359, 0.0003, 0.0036},
		{-0.0024, -0.0142, 0.0214 /*...,*/, -0.0036, 0.0042, 0.0056},
		{0.0172, 0.1405, 0.0023 /*...,*/, -0.0057, -0.0426, 0.0021},
		/*...,*/
		{-0.0394, -0.0555, 0.0665 /*...,*/, -0.0465, -0.0142, 0.0180},
		{-0.0068, -0.0230, -0.0009 /*...,*/, -0.0149, -0.0061, 0.0051},
		{-0.0181, 0.0975, -0.0097 /*...,*/, 0.0082, 0.0026, 0.0075},
	}

	actualAttnNormPart, err := transformerBlock.attn_norm.doNormalization(x)
	if err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedAttnNormPart, expectedAttnNormPartSize, actualAttnNormPart, 2*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	actualAttnNormalizedX, err := ml.MultiplyElementwise(actualAttnNormPart, transformerBlock.attn_norm.weights)
	if err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedAttnNormalizedX, expectedAttnNormalizedXSize, actualAttnNormalizedX, common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}
	return actualAttnNormalizedX
}

func testTransformerBlock_Attention_Forward(t *testing.T, skipCompareTestTensor bool, infContext *InferenceContext, attention *LlamaAttention, x *ml.Tensor, startPos int, freqsCis *ml.Tensor, mask *ml.Tensor) *ml.Tensor {
	expectedXqSize := []int{15, 4096}
	expectedXq := [][]float32{
		{-0.0974, -0.1632, 0.2617 /*...,*/, 0.4542, -0.3723, 0.2629},
		{0.0206, -0.0089, 0.1009 /*...,*/, 0.0657, -0.1027, 0.1156},
		{0.1101, -0.1451, 0.3900 /*...,*/, 2.1700, -3.1855, 2.1169},
		/*...,*/
		{0.3651, -0.0470, 0.8779 /*...,*/, 1.7904, -2.7587, 2.0914},
		{0.0360, -0.0227, 0.0595 /*...,*/, 0.0323, -0.0557, 0.0191},
		{0.4571, -0.1180, 1.4996 /*...,*/, 2.3494, -2.4556, 1.4067},
	}

	expectedXkSize := []int{15, 1024}
	expectedXk := [][]float32{
		{0.5587, -0.3365, 0.9848 /*...,*/, 1.7037, -2.0537, 1.7776},
		{0.4686, 0.1622, 0.0747 /*...,*/, -0.0822, 0.0329, 0.0459},
		{7.8971, 3.4503, 1.1751 /*...,*/, -1.5299, 1.5073, -0.8149},
		/*...,*/
		{7.6970, 3.1144, 2.2046 /*...,*/, -1.2853, 1.7026, -0.3482},
		{0.3218, 0.2011, 0.0825 /*...,*/, -0.0980, 0.0964, 0.0129},
		{2.9652, -0.4350, 2.7118 /*...,*/, 0.3060, -0.6938, -2.2576},
	}

	expectedXvSize := []int{15, 1024}
	expectedXv := [][]float32{
		{0.0109, -0.0005, 0.0286 /*...,*/, 0.0007, -0.0010, 0.0004},
		{-0.0047, -0.0017, -0.0118 /*...,*/, 0.0064, -0.0113, 0.0087},
		{0.1022, -0.0047, 0.0568 /*...,*/, 0.0708, 0.0168, -0.0362},
		/*...,*/
		{0.0217, -0.0317, 0.0653 /*...,*/, 0.0031, -0.0533, 0.0256},
		{-0.0027, 0.0004, -0.0056 /*...,*/, 0.0035, -0.0044, 0.0066},
		{0.0008, -0.0022, 0.0154 /*...,*/, 0.0167, -0.0084, -0.0090},
	}

	sequenceLength := x.Size[0]

	// lat.attn_wq: [out_features, in_features] -> shape: [4096 4096] -> [N_Heads * HeadDim, Dim]
	actualXq, err := ml.LinearTransformation(x, attention.attn_wq)
	if err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXq, expectedXqSize, actualXq, 4*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	// lat.attn_wk: [out_features, in_features] -> shape: [1024 4096] -> [N_KVHeads * HeadDim, Dim]
	actualXk, err := ml.LinearTransformation(x, attention.attn_wk)
	if err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXk, expectedXkSize, actualXk, 8*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	// lat.attn_wv: [out_features, in_features] -> shape: [1024 4096] -> [N_KVHeads * HeadDim, Dim]
	actualXv, err := ml.LinearTransformation(x, attention.attn_wv)
	if err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXv, expectedXvSize, actualXv, 4*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		Do reshapings
	*/

	expectedXqRsSize := []int{15, 32, 128}
	expectedXqRs := [][][]float32{
		{
			{-9.7422e-02, -1.6317e-01, 2.6171e-01 /*...,*/, 8.7456e-01, 4.5385e-01, 3.1548e-01},
			{1.3232e-01, -4.4514e-02, 3.5850e-01 /*...,*/, 5.4703e-01, 3.7532e-01, 3.4740e-01},
			{6.9574e-01, -4.4043e-01, 9.2948e-01 /*...,*/, 9.5450e-01, 7.0820e-01, 4.0560e-01},
			/*...,*/
			{2.5047e-01, -5.5406e-01, 1.2645e+00 /*...,*/, 6.2868e-01, -6.0493e-02, 1.5389e+00},
			{-2.6671e-01, -3.8844e-01, -3.0750e-01 /*...,*/, 4.5032e-01, -4.7121e-01, -1.3882e-01},
			{-5.6102e-01, 3.7558e-01, 2.1194e-01 /*...,*/, 4.5418e-01, -3.7228e-01, 2.6291e-01},
		},
		{
			{2.0646e-02, -8.8847e-03, 1.0093e-01 /*...,*/, 1.2186e-01, 6.7169e-02, 5.3663e-02},
			{-3.0246e-02, 1.8086e-03, -3.1814e-02 /*...,*/, 8.5876e-02, 1.1263e-01, 1.0839e-01},
			{1.8133e-01, -1.2319e-01, 2.5101e-01 /*...,*/, 1.5047e-01, 1.7557e-01, 1.7040e-01},
			/*...,*/
			{-7.3076e-02, -1.5629e-01, -3.2961e-02 /*...,*/, 2.1952e-02, -6.7880e-02, 2.5287e-02},
			{1.1608e-01, 1.5886e-02, 2.8677e-02 /*...,*/, -1.3075e-02, -6.4812e-02, -1.7991e-02},
			{-8.8719e-02, -6.0895e-03, 2.4713e-02 /*...,*/, 6.5737e-02, -1.0265e-01, 1.1557e-01},
		},
		{
			{1.1011e-01, -1.4510e-01, 3.9000e-01 /*...,*/, 2.3012e+00, 1.2982e+00, 1.4781e+00},
			{2.4087e-01, 4.4564e-02, 1.2808e-01 /*...,*/, 1.8443e+00, 1.9174e+00, 1.6965e+00},
			{1.8177e+00, -7.8401e-01, 2.5387e+00 /*...,*/, 2.0767e+00, 2.4926e+00, 2.3336e+00},
			/*...,*/
			{-3.1435e-01, -1.9310e+00, 4.2275e-01 /*...,*/, 1.2577e+00, -2.0807e+00, -1.7765e-01},
			{2.0104e+00, 3.0268e-01, -3.2210e-02 /*...,*/, 8.9383e-01, -1.1510e+00, 1.2941e-01},
			{-8.9192e-01, -1.0077e+00, 8.7606e-02 /*...,*/, 2.1700e+00, -3.1855e+00, 2.1169e+00},
		},
		/*...,*/
		{
			{3.6513e-01, -4.6951e-02, 8.7788e-01 /*...,*/, 2.2966e+00, 1.4644e+00, 1.2319e+00},
			{-1.8565e-01, 7.5481e-02, -1.9541e-01 /*...,*/, 1.5944e+00, 1.6606e+00, 1.7784e+00},
			{2.9575e+00, -1.7174e+00, 3.0296e+00 /*...,*/, 2.0744e+00, 2.6284e+00, 2.3956e+00},
			/*...,*/
			{-1.0103e+00, -1.7282e+00, 5.3124e-01 /*...,*/, 9.4434e-01, -2.0660e+00, -1.9275e-01},
			{2.2196e+00, 5.6332e-03, -2.7980e-01 /*...,*/, 8.1934e-01, -9.2910e-01, 5.6069e-02},
			{-4.7546e-01, -5.8105e-01, 3.4718e-02 /*...,*/, 1.7904e+00, -2.7587e+00, 2.0914e+00},
		},
		{
			{3.5971e-02, -2.2706e-02, 5.9504e-02 /*...,*/, 6.1715e-02, 4.8816e-02, 2.9821e-02},
			{9.8976e-02, -6.8781e-02, 5.9712e-02 /*...,*/, 4.3594e-02, 7.4092e-02, 8.5981e-02},
			{1.9424e-01, -1.2670e-01, 2.0396e-01 /*...,*/, 1.1370e-01, 1.2174e-01, 1.1693e-01},
			/*...,*/
			{-5.6662e-03, -8.3996e-02, -8.1616e-02 /*...,*/, 4.0625e-02, -8.5604e-02, -2.7477e-02},
			{4.6506e-02, 5.5263e-03, 4.2742e-03 /*...,*/, 1.5699e-02, -3.0517e-02, -1.5159e-02},
			{-4.6612e-02, 4.6806e-03, 2.8639e-02 /*...,*/, 3.2324e-02, -5.5728e-02, 1.9095e-02},
		},
		{
			{4.5710e-01, -1.1800e-01, 1.4996e+00 /*...,*/, 2.6436e+00, 1.2102e+00, 5.0838e-01},
			{9.6557e-02, -3.4989e-02, 1.2897e-01 /*...,*/, 2.2276e+00, 1.2230e+00, 1.0092e+00},
			{2.5894e+00, -1.1153e+00, 2.7558e+00 /*...,*/, 2.9313e+00, 2.2336e+00, 1.2770e+00},
			/*...,*/
			{2.4838e-02, -3.2833e+00, 3.2588e+00 /*...,*/, 3.2141e+00, -1.1062e+00, 7.0349e-01},
			{3.0594e+00, -1.6547e+00, -1.3107e+00 /*...,*/, 1.1335e+00, -1.1497e+00, 6.4756e-02},
			{2.7660e-01, -2.6660e-01, -3.4703e-02 /*...,*/, 2.3494e+00, -2.4556e+00, 1.4067e+00},
		},
	}

	expectedXkRsSize := []int{15, 8, 128}
	expectedXkRs := [][][]float32{
		{
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{1.7906e-02, -1.0323e-02, -7.5300e-02 /*...,*/, 8.1819e-01, -5.2369e-01, 6.0337e-02},
			{6.6542e-02, -3.7643e-02, -5.2828e-02 /*...,*/, 1.8952e+00, -1.0149e+00, -9.9316e-01},
			/*...,*/
			{-1.9670e-01, 1.1818e-01, -1.3081e-01 /*...,*/, -7.6494e-01, 1.5089e+00, 7.6288e-01},
			{-4.9350e-02, 9.3404e-03, -8.0326e-02 /*...,*/, -4.9850e-01, 2.8490e-01, 3.8502e-01},
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
		},
		{
			{4.6858e-01, 1.6215e-01, 7.4740e-02 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{-5.2816e-02, -6.3121e-02, -4.6727e-02 /*...,*/, -7.4227e-02, -4.0639e-02, -1.3300e-01},
			{2.1797e-02, -4.1893e-02, -2.9867e-02 /*...,*/, -1.3752e-01, -4.9149e-02, 9.6101e-02},
			/*...,*/
			{-1.1077e-01, -3.8428e-02, 9.4273e-02 /*...,*/, 3.1714e-02, 4.0120e-02, -1.0732e-01},
			{-1.6118e-01, 1.1211e-01, -2.0423e-01 /*...,*/, 1.4500e-01, -1.1386e-01, -8.8506e-02},
			{-9.9322e-03, -3.0433e-01, -1.2610e-01 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
		},
		{
			{7.8971e+00, 3.4503e+00, 1.1751e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			{-4.2669e-01, -3.2506e+00, -2.2153e+00 /*...,*/, -1.5790e+00, 4.3659e-01, -1.3140e+00},
			{5.1532e-01, 5.3257e-01, -6.1076e-01 /*...,*/, -1.0003e+00, 8.3557e-01, 4.1634e-01},
			/*...,*/
			{-3.1664e+00, -1.1471e+00, 1.6622e+00 /*...,*/, 1.6303e+00, -1.5842e+00, -5.9981e-01},
			{-4.3509e+00, 7.4818e-01, -3.8100e+00 /*...,*/, 3.6117e+00, -2.1949e+00, -2.0049e+00},
			{6.6238e-01, -5.2133e+00, -3.5139e+00 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
		},
		/*...,*/
		{
			{7.6970e+00, 3.1144e+00, 2.2046e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{-4.0040e-01, -3.5385e+00, -1.6972e+00 /*...,*/, -1.9167e+00, 8.1168e-01, -9.3656e-01},
			{5.0167e-02, 1.4137e+00, -3.2782e-01 /*...,*/, -1.8262e+00, 1.8397e+00, -9.6362e-02},
			/*...,*/
			{-2.7194e+00, 2.1670e-02, 9.4138e-01 /*...,*/, 1.6307e+00, -1.6633e+00, -7.0559e-01},
			{-4.8732e+00, -6.3733e-01, -3.9122e+00 /*...,*/, 2.2310e+00, -4.1900e+00, -2.0448e+00},
			{1.0723e+00, -5.1179e+00, -3.6607e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4816e-01},
		},
		{
			{3.2184e-01, 2.0114e-01, 8.2546e-02 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{-6.0029e-03, -4.6157e-02, -3.6131e-02 /*...,*/, -4.9256e-02, 4.1981e-02, -7.7876e-02},
			{-8.0149e-02, 1.3749e-01, 5.9112e-02 /*...,*/, 4.3934e-02, -1.3365e-01, 1.3886e-01},
			/*...,*/
			{-6.8796e-02, -1.1835e-02, 4.8842e-02 /*...,*/, 6.2568e-02, 3.0143e-02, -9.3066e-02},
			{-1.7158e-01, 4.0830e-02, -1.9338e-01 /*...,*/, 2.5320e-02, -8.7856e-02, -7.2434e-02},
			{2.3165e-02, -1.8819e-01, -2.5320e-02 /*...,*/, -9.7979e-02, 9.6370e-02, 1.2926e-02},
		},
		{
			{2.9652e+00, -4.3503e-01, 2.7118e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
			{-1.2986e-01, -1.3757e-01, -6.7788e-02 /*...,*/, -9.4586e-01, 2.0041e-02, -4.7395e-01},
			{-2.3858e+00, 5.7992e+00, 3.3556e+00 /*...,*/, -3.5978e-02, -1.6827e-01, 1.8953e+00},
			/*...,*/
			{-1.7856e+00, 1.0335e+00, 2.3853e+00 /*...,*/, 2.4644e+00, 2.8289e-01, -4.2407e+00},
			{-3.8269e-01, 3.4285e-01, -1.0427e+00 /*...,*/, 8.0458e-01, -1.1823e+00, -9.9586e-01},
			{8.6834e-02, -1.5074e-01, -1.9138e+00 /*...,*/, 3.0597e-01, -6.9376e-01, -2.2576e+00},
		},
	}

	expectedXvRsSize := []int{15, 8, 128}
	expectedXvRs := [][][]float32{
		{
			{1.0913e-02, -4.7801e-04, 2.8640e-02 /*...,*/, -2.1515e-02, -2.6500e-03, 4.7764e-02},
			{1.6308e-04, 1.4786e-03, -5.2374e-04 /*...,*/, 4.0014e-04, -4.4112e-04, -9.5709e-04},
			{1.0466e-03, -2.5886e-03, 6.4222e-04 /*...,*/, -2.1045e-03, 1.5630e-03, 1.1579e-03},
			/*...,*/
			{6.1011e-04, 2.3969e-04, -7.9793e-04 /*...,*/, 1.9225e-03, 7.9390e-04, 9.3130e-04},
			{4.0138e-04, 3.6933e-03, -3.3550e-04 /*...,*/, -2.8911e-04, 1.1221e-03, -2.7706e-03},
			{2.4258e-03, 1.1657e-03, 7.0518e-04 /*...,*/, 6.6461e-04, -9.7815e-04, 3.5757e-04},
		},
		{
			{-4.7307e-03, -1.7014e-03, -1.1822e-02 /*...,*/, -8.5324e-03, 1.0732e-02, -2.6448e-04},
			{-1.3907e-03, -1.5558e-03, -1.3331e-03 /*...,*/, -9.1358e-03, -7.0879e-04, -1.3888e-03},
			{-3.6529e-03, 9.9129e-03, 2.3944e-03 /*...,*/, -1.1497e-02, -1.0619e-02, -2.0509e-03},
			/*...,*/
			{-3.2024e-03, 5.1106e-03, -6.0006e-05 /*...,*/, 1.5563e-03, 1.2662e-03, 3.2343e-04},
			{-5.1777e-03, 3.0121e-03, 1.1942e-03 /*...,*/, -1.0517e-03, 1.0548e-02, -4.6303e-03},
			{-1.7185e-04, 2.8876e-03, -1.3695e-02 /*...,*/, 6.3693e-03, -1.1288e-02, 8.6698e-03},
		},
		{
			{1.0218e-01, -4.7389e-03, 5.6817e-02 /*...,*/, -1.5799e-02, -1.2324e-02, -9.1933e-02},
			{3.8830e-02, -3.8196e-02, 2.2808e-02 /*...,*/, -2.8040e-02, -5.0615e-02, -7.5970e-02},
			{6.4443e-02, 1.7406e-02, -5.4531e-03 /*...,*/, -5.5268e-02, -1.8770e-02, -2.8677e-02},
			/*...,*/
			{-3.2645e-02, -3.3754e-02, -1.2851e-02 /*...,*/, -3.6677e-02, 2.1962e-03, 2.7318e-02},
			{-1.1419e-02, 1.3481e-02, -2.9121e-02 /*...,*/, -2.8687e-02, -1.6453e-02, 3.0407e-03},
			{-2.4622e-02, -6.0911e-02, 1.6443e-02 /*...,*/, 7.0776e-02, 1.6823e-02, -3.6186e-02},
		},
		/*...,*/
		{
			{2.1686e-02, -3.1682e-02, 6.5300e-02 /*...,*/, 2.0801e-02, 4.5545e-02, -6.0511e-02},
			{5.5816e-02, -7.5619e-02, -1.0595e-02 /*...,*/, 9.9449e-02, -6.5036e-02, -2.9054e-02},
			{2.6425e-02, -7.3392e-03, -1.5957e-02 /*...,*/, 3.2710e-02, -2.3600e-02, -4.9277e-02},
			/*...,*/
			{-2.5396e-02, -1.6222e-02, -4.6781e-02 /*...,*/, 5.1221e-02, 7.1369e-03, -2.3695e-02},
			{5.1342e-02, -4.0454e-02, 2.7118e-03 /*...,*/, 8.5404e-02, 1.0772e-02, 2.2609e-02},
			{-3.8587e-02, -1.0937e-02, -1.8124e-02 /*...,*/, 3.1006e-03, -5.3299e-02, 2.5574e-02},
		},
		{
			{-2.7372e-03, 3.8026e-04, -5.5800e-03 /*...,*/, -5.2489e-03, 4.3830e-03, -1.1022e-02},
			{-1.6923e-03, 1.1716e-03, 1.4530e-04 /*...,*/, 1.1970e-03, 4.3336e-03, -1.0145e-04},
			{-3.3312e-03, 2.6650e-03, 3.6644e-04 /*...,*/, -2.2175e-03, -9.5940e-03, -5.6057e-03},
			/*...,*/
			{-9.2908e-04, 1.4751e-03, -3.8694e-03 /*...,*/, 2.0414e-03, -3.7970e-03, 7.0636e-04},
			{-2.8142e-03, -1.4741e-03, 3.4434e-03 /*...,*/, 1.0130e-03, 1.4796e-03, -4.4738e-03},
			{-1.1751e-02, 1.4951e-03, 7.8496e-03 /*...,*/, 3.4656e-03, -4.3663e-03, 6.5875e-03},
		},

		{
			{7.5373e-04, -2.2473e-03, 1.5444e-02 /*...,*/, -3.6623e-03, 3.7663e-03, 1.2237e-01},
			{-1.9953e-03, 6.8435e-03, 2.8523e-03 /*...,*/, -2.0271e-03, -3.7172e-03, 8.0297e-03},
			{1.8759e-02, -9.9398e-03, 8.4348e-03 /*...,*/, 7.5021e-03, 2.6062e-02, -3.6960e-03},
			/*...,*/
			{-4.0907e-03, 1.0375e-02, -1.3782e-02 /*...,*/, 1.2689e-02, 1.7571e-02, 3.1714e-02},
			{1.0694e-03, -1.9459e-03, -5.4336e-03 /*...,*/, 3.8816e-03, 7.1745e-03, 9.8520e-03},
			{-9.6662e-03, -1.1503e-02, 1.5865e-02 /*...,*/, 1.6660e-02, -8.3912e-03, -9.0059e-03},
		},
	}

	if actualXq, err = actualXq.Reshape([]int{sequenceLength, attention.N_Heads, attention.HeadDim}); err != nil {
		t.Fatal(err)
	}

	if actualXk, err = actualXk.Reshape([]int{sequenceLength, attention.N_KVHeads, attention.HeadDim}); err != nil {
		t.Fatal(err)
	}

	if actualXv, err = actualXv.Reshape([]int{sequenceLength, attention.N_KVHeads, attention.HeadDim}); err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXqRs, expectedXqRsSize, actualXq, 4*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXkRs, expectedXkRsSize, actualXk, 8*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXvRs, expectedXvRsSize, actualXv, 4*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		Apply rotary embeddings
	*/

	expectedXqRotarySize := []int{15, 32, 128}
	expectedXqRotary := [][][]float32{
		{
			{-9.7422e-02, -1.6317e-01, 2.6171e-01 /*...,*/, 8.7456e-01, 4.5385e-01, 3.1548e-01},
			{1.3232e-01, -4.4514e-02, 3.5850e-01 /*...,*/, 5.4703e-01, 3.7532e-01, 3.4740e-01},
			{6.9574e-01, -4.4043e-01, 9.2948e-01 /*...,*/, 9.5450e-01, 7.0820e-01, 4.0560e-01},
			/*...,*/
			{2.5047e-01, -5.5406e-01, 1.2645e+00 /*...,*/, 6.2868e-01, -6.0493e-02, 1.5389e+00},
			{-2.6671e-01, -3.8844e-01, -3.0750e-01 /*...,*/, 4.5032e-01, -4.7121e-01, -1.3882e-01},
			{-5.6102e-01, 3.7558e-01, 2.1194e-01 /*...,*/, 4.5418e-01, -3.7228e-01, 2.6291e-01},
		},
		{
			{1.8631e-02, 1.2573e-02, 1.2035e-01 /*...,*/, 1.2186e-01, 6.7169e-02, 5.3663e-02},
			{-1.7864e-02, -2.4474e-02, -2.5341e-02 /*...,*/, 8.5876e-02, 1.1263e-01, 1.0839e-01},
			{2.0163e-01, 8.6028e-02, 2.1220e-01 /*...,*/, 1.5047e-01, 1.7557e-01, 1.7040e-01},
			/*...,*/
			{9.2031e-02, -1.4594e-01, -2.4383e-02 /*...,*/, 2.1952e-02, -6.7880e-02, 2.5287e-02},
			{4.9349e-02, 1.0626e-01, 2.8606e-02 /*...,*/, -1.3075e-02, -6.4812e-02, -1.7991e-02},
			{-4.2811e-02, -7.7944e-02, -1.7524e-02 /*...,*/, 6.5737e-02, -1.0265e-01, 1.1557e-01},
		},
		{
			{8.6118e-02, 1.6051e-01, 3.6869e-01 /*...,*/, 2.3012e+00, 1.2982e+00, 1.4781e+00},
			{-1.4076e-01, 2.0047e-01, -1.6284e-01 /*...,*/, 1.8443e+00, 1.9174e+00, 1.6965e+00},
			{-4.3524e-02, 1.9791e+00, 7.8367e-01 /*...,*/, 2.0767e+00, 2.4926e+00, 2.3336e+00},
			/*...,*/
			{1.8867e+00, 5.1775e-01, -3.5486e-01 /*...,*/, 1.2577e+00, -2.0807e+00, -1.7765e-01},
			{-1.1119e+00, 1.7021e+00, 2.9340e-01 /*...,*/, 8.9383e-01, -1.1510e+00, 1.2941e-01},
			{1.2875e+00, -3.9165e-01, -4.0549e-01 /*...,*/, 2.1700e+00, -3.1855e+00, 2.1169e+00},
		},
		/*...,*/
		{
			{2.8292e-01, -2.3554e-01, -1.0004e+00 /*...,*/, 2.2966e+00, 1.4643e+00, 1.2319e+00},
			{-1.1616e-01, 1.6331e-01, 1.9293e-01 /*...,*/, 1.5944e+00, 1.6606e+00, 1.7784e+00},
			{1.5742e+00, -3.0362e+00, -3.0963e+00 /*...,*/, 2.0745e+00, 2.6284e+00, 2.3956e+00},
			/*...,*/
			{-1.7799e+00, -9.1624e-01, -3.7893e-01 /*...,*/, 9.4433e-01, -2.0660e+00, -1.9276e-01},
			{1.8760e+00, -1.1862e+00, 1.8114e-01 /*...,*/, 8.1934e-01, -9.2910e-01, 5.6066e-02},
			{-7.1299e-01, -2.3520e-01, 3.5515e-02 /*...,*/, 1.7904e+00, -2.7587e+00, 2.0914e+00},
		},
		{
			{4.2182e-02, -5.4906e-03, -6.5568e-02 /*...,*/, 6.1715e-02, 4.8816e-02, 2.9821e-02},
			{1.1872e-01, -2.0828e-02, -2.8896e-02 /*...,*/, 4.3594e-02, 7.4091e-02, 8.5981e-02},
			{2.2950e-01, -3.3358e-02, -9.5946e-02 /*...,*/, 1.1370e-01, 1.2174e-01, 1.1693e-01},
			/*...,*/
			{3.0151e-02, -7.8603e-02, 2.5738e-03 /*...,*/, 4.0625e-02, -8.5604e-02, -2.7478e-02},
			{3.9879e-02, 2.4555e-02, -2.2776e-02 /*...,*/, 1.5699e-02, -3.0517e-02, -1.5159e-02},
			{-4.4265e-02, -1.5338e-02, 5.6960e-03 /*...,*/, 3.2324e-02, -5.5728e-02, 1.9095e-02},
		},
		{
			{1.7939e-01, 4.3667e-01, -1.2882e-01 /*...,*/, 2.6436e+00, 1.2102e+00, 5.0839e-01},
			{4.7863e-02, 9.0866e-02, 1.3867e-01 /*...,*/, 2.2276e+00, 1.2230e+00, 1.0092e+00},
			{1.4589e+00, 2.4126e+00, 6.4037e-01 /*...,*/, 2.9313e+00, 2.2336e+00, 1.2770e+00},
			/*...,*/
			{3.2559e+00, -4.2435e-01, 2.9481e+00 /*...,*/, 3.2141e+00, -1.1062e+00, 7.0349e-01},
			{2.0575e+00, 2.8044e+00, -4.2136e-01 /*...,*/, 1.1335e+00, -1.1497e+00, 6.4751e-02},
			{3.0192e-01, 2.3755e-01, 6.2483e-01 /*...,*/, 2.3494e+00, -2.4556e+00, 1.4067e+00},
		},
	}

	expectedXkRotarySize := []int{15, 8, 128}
	expectedXkRotary := [][][]float32{
		{
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{1.7906e-02, -1.0323e-02, -7.5300e-02 /*...,*/, 8.1819e-01, -5.2369e-01, 6.0337e-02},
			{6.6542e-02, -3.7643e-02, -5.2828e-02 /*...,*/, 1.8952e+00, -1.0149e+00, -9.9316e-01},
			/*...,*/
			{-1.9670e-01, 1.1818e-01, -1.3081e-01 /*...,*/, -7.6494e-01, 1.5089e+00, 7.6288e-01},
			{-4.9350e-02, 9.3404e-03, -8.0326e-02 /*...,*/, -4.9850e-01, 2.8490e-01, 3.8502e-01},
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
		},
		{
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{2.4578e-02, -7.8547e-02, -1.1616e-01 /*...,*/, -7.4227e-02, -4.0639e-02, -1.3300e-01},
			{4.7029e-02, -4.2932e-03, -3.3278e-02 /*...,*/, -1.3752e-01, -4.9149e-02, 9.6101e-02},
			/*...,*/
			{-2.7512e-02, -1.1397e-01, 1.2006e-01 /*...,*/, 3.1714e-02, 4.0120e-02, -1.0732e-01},
			{-1.8142e-01, -7.5052e-02, 4.6226e-03 /*...,*/, 1.4500e-01, -1.1386e-01, -8.8506e-02},
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
		},
		{
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			{3.1334e+00, 9.6475e-01, -3.5259e+00 /*...,*/, -1.5790e+00, 4.3659e-01, -1.3140e+00},
			{-6.9871e-01, 2.4695e-01, -5.1941e-01 /*...,*/, -1.0003e+00, 8.3557e-01, 4.1634e-01},
			/*...,*/
			{2.3607e+00, -2.4019e+00, 9.4849e-01 /*...,*/, 1.6303e+00, -1.5842e+00, -5.9981e-01},
			{1.1303e+00, -4.2676e+00, 3.3445e+00 /*...,*/, 3.6117e+00, -2.1949e+00, -2.0049e+00},
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
		},
		/*...,*/
		{
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{-2.2365e+00, -2.7711e+00, 2.7901e+00 /*...,*/, -1.9167e+00, 8.1169e-01, -9.3656e-01},
			{8.0089e-01, 1.1660e+00, 3.4041e-01 /*...,*/, -1.8262e+00, 1.8397e+00, -9.6356e-02},
			/*...,*/
			{-2.2832e+00, 1.4774e+00, -1.2831e+00 /*...,*/, 1.6308e+00, -1.6633e+00, -7.0560e-01},
			{-4.4542e+00, 2.0770e+00, 2.5727e+00 /*...,*/, 2.2310e+00, -4.1900e+00, -2.0448e+00},
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
		},
		{
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{1.3946e-02, -4.4407e-02, 7.1611e-02 /*...,*/, -4.9256e-02, 4.1981e-02, -7.7875e-02},
			{-1.3050e-01, 9.1093e-02, 4.4730e-02 /*...,*/, 4.3934e-02, -1.3365e-01, 1.3886e-01},
			/*...,*/
			{-5.7456e-02, -3.9645e-02, -6.8386e-02 /*...,*/, 6.2568e-02, 3.0144e-02, -9.3066e-02},
			{-1.7285e-01, -3.5040e-02, -1.7397e-02 /*...,*/, 2.5321e-02, -8.7856e-02, -7.2434e-02},
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
		},
		{
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
			{1.1852e-01, -1.4745e-01, 5.5575e-01 /*...,*/, -9.4587e-01, 2.0043e-02, -4.7395e-01},
			{-6.0710e+00, -1.5705e+00, 1.9432e+00 /*...,*/, -3.5979e-02, -1.6828e-01, 1.8953e+00},
			/*...,*/
			{-1.2680e+00, -1.6275e+00, 5.2333e-01 /*...,*/, 2.4644e+00, 2.8291e-01, -4.2407e+00},
			{-3.9196e-01, -3.3221e-01, -3.9067e-02 /*...,*/, 8.0458e-01, -1.1823e+00, -9.9587e-01},
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
		},
	}

	if actualXq, actualXk, err = applyRotaryEmbeddings(actualXq, actualXk, freqsCis); err != nil {
		return nil
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXqRotary, expectedXqRotarySize, actualXq, 9*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXkRotary, expectedXkRotarySize, actualXk, 12*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		Update KV cache
	*/

	infContext.CacheK[attention.LayerIndex].SetSlice([]int{startPos}, []int{startPos + sequenceLength}, actualXk)
	infContext.CacheV[attention.LayerIndex].SetSlice([]int{startPos}, []int{startPos + sequenceLength}, actualXv)

	/*
		Retrieve cached KV so far
	*/

	actualKeys, err := infContext.CacheK[attention.LayerIndex].Slice([]int{0}, []int{startPos + sequenceLength})
	if err != nil {
		t.Fatal(err)
	}
	actualValues, err := infContext.CacheV[attention.LayerIndex].Slice([]int{0}, []int{startPos + sequenceLength})
	if err != nil {
		t.Fatal(err)
	}

	/*
		Repeat k/v heads if N_KVHeads < N_Heads
	*/

	expectedKeysRepSize := []int{15, 32, 128}
	expectedKeysRep := [][][]float32{
		{
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			/*...,*/
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
		},
		{
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			/*...,*/
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
		},
		{
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			/*...,*/
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
		},
		/*...,*/
		{
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			/*...,*/
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
		},
		{
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			/*...,*/
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
		},
		{
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
			/*...,*/
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
		},
	}

	expectedValuesRepSize := []int{15, 32, 128}
	expectedValuesRep := [][][]float32{
		{
			{0.0109, -0.0005, 0.0286 /*...,*/, -0.0215, -0.0026, 0.0478},
			{0.0109, -0.0005, 0.0286 /*...,*/, -0.0215, -0.0026, 0.0478},
			{0.0109, -0.0005, 0.0286 /*...,*/, -0.0215, -0.0026, 0.0478},
			/*...,*/
			{0.0024, 0.0012, 0.0007 /*...,*/, 0.0007, -0.0010, 0.0004},
			{0.0024, 0.0012, 0.0007 /*...,*/, 0.0007, -0.0010, 0.0004},
			{0.0024, 0.0012, 0.0007 /*...,*/, 0.0007, -0.0010, 0.0004},
		},
		{
			{-0.0047, -0.0017, -0.0118 /*...,*/, -0.0085, 0.0107, -0.0003},
			{-0.0047, -0.0017, -0.0118 /*...,*/, -0.0085, 0.0107, -0.0003},
			{-0.0047, -0.0017, -0.0118 /*...,*/, -0.0085, 0.0107, -0.0003},
			/*...,*/
			{-0.0002, 0.0029, -0.0137 /*...,*/, 0.0064, -0.0113, 0.0087},
			{-0.0002, 0.0029, -0.0137 /*...,*/, 0.0064, -0.0113, 0.0087},
			{-0.0002, 0.0029, -0.0137 /*...,*/, 0.0064, -0.0113, 0.0087},
		},
		{
			{0.1022, -0.0047, 0.0568 /*...,*/, -0.0158, -0.0123, -0.0919},
			{0.1022, -0.0047, 0.0568 /*...,*/, -0.0158, -0.0123, -0.0919},
			{0.1022, -0.0047, 0.0568 /*...,*/, -0.0158, -0.0123, -0.0919},
			/*...,*/
			{-0.0246, -0.0609, 0.0164 /*...,*/, 0.0708, 0.0168, -0.0362},
			{-0.0246, -0.0609, 0.0164 /*...,*/, 0.0708, 0.0168, -0.0362},
			{-0.0246, -0.0609, 0.0164 /*...,*/, 0.0708, 0.0168, -0.0362},
		},
		/*...,*/
		{
			{0.0217, -0.0317, 0.0653 /*...,*/, 0.0208, 0.0455, -0.0605},
			{0.0217, -0.0317, 0.0653 /*...,*/, 0.0208, 0.0455, -0.0605},
			{0.0217, -0.0317, 0.0653 /*...,*/, 0.0208, 0.0455, -0.0605},
			/*...,*/
			{-0.0386, -0.0109, -0.0181 /*...,*/, 0.0031, -0.0533, 0.0256},
			{-0.0386, -0.0109, -0.0181 /*...,*/, 0.0031, -0.0533, 0.0256},
			{-0.0386, -0.0109, -0.0181 /*...,*/, 0.0031, -0.0533, 0.0256},
		},
		{
			{-0.0027, 0.0004, -0.0056 /*...,*/, -0.0052, 0.0044, -0.0110},
			{-0.0027, 0.0004, -0.0056 /*...,*/, -0.0052, 0.0044, -0.0110},
			{-0.0027, 0.0004, -0.0056 /*...,*/, -0.0052, 0.0044, -0.0110},
			/*...,*/
			{-0.0118, 0.0015, 0.0078 /*...,*/, 0.0035, -0.0044, 0.0066},
			{-0.0118, 0.0015, 0.0078 /*...,*/, 0.0035, -0.0044, 0.0066},
			{-0.0118, 0.0015, 0.0078 /*...,*/, 0.0035, -0.0044, 0.0066},
		},
		{
			{0.0008, -0.0022, 0.0154 /*...,*/, -0.0037, 0.0038, 0.1224},
			{0.0008, -0.0022, 0.0154 /*...,*/, -0.0037, 0.0038, 0.1224},
			{0.0008, -0.0022, 0.0154 /*...,*/, -0.0037, 0.0038, 0.1224},
			/*...,*/
			{-0.0097, -0.0115, 0.0159 /*...,*/, 0.0167, -0.0084, -0.0090},
			{-0.0097, -0.0115, 0.0159 /*...,*/, 0.0167, -0.0084, -0.0090},
			{-0.0097, -0.0115, 0.0159 /*...,*/, 0.0167, -0.0084, -0.0090},
		},
	}

	N_Rep := attention.N_Rep

	if actualKeys, err = attentionRepeatKV(actualKeys, N_Rep); err != nil { // shape=[5, 32, 128] (cacheLen + sequenceLength, N_Heads, HeadDim)
		t.Fatal(err)
	}
	if actualValues, err = attentionRepeatKV(actualValues, N_Rep); err != nil { // shape=[5, 32, 128] (cacheLen + sequenceLength, N_Heads, HeadDim)
		t.Fatal(err)
	}

	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedKeysRep, expectedKeysRepSize, actualKeys, 12*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedValuesRep, expectedValuesRepSize, actualValues, 7*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		Do transposes
	*/

	expectedXqTransposeSize := []int{32, 15, 128}
	expectedXqTranspose := [][][]float32{
		{
			{-9.7422e-02, -1.6317e-01, 2.6171e-01 /*...,*/, 8.7456e-01, 4.5385e-01, 3.1548e-01},
			{1.8631e-02, 1.2573e-02, 1.2035e-01 /*...,*/, 1.2186e-01, 6.7169e-02, 5.3663e-02},
			{8.6118e-02, 1.6051e-01, 3.6869e-01 /*...,*/, 2.3012e+00, 1.2982e+00, 1.4781e+00},
			/*...,*/
			{2.8292e-01, -2.3554e-01, -1.0004e+00 /*...,*/, 2.2966e+00, 1.4643e+00, 1.2319e+00},
			{4.2182e-02, -5.4906e-03, -6.5568e-02 /*...,*/, 6.1715e-02, 4.8816e-02, 2.9821e-02},
			{1.7939e-01, 4.3667e-01, -1.2882e-01 /*...,*/, 2.6436e+00, 1.2102e+00, 5.0839e-01},
		},
		{
			{1.3232e-01, -4.4514e-02, 3.5850e-01 /*...,*/, 5.4703e-01, 3.7532e-01, 3.4740e-01},
			{-1.7864e-02, -2.4474e-02, -2.5341e-02 /*...,*/, 8.5876e-02, 1.1263e-01, 1.0839e-01},
			{-1.4076e-01, 2.0047e-01, -1.6284e-01 /*...,*/, 1.8443e+00, 1.9174e+00, 1.6965e+00},
			/*...,*/
			{-1.1616e-01, 1.6331e-01, 1.9293e-01 /*...,*/, 1.5944e+00, 1.6606e+00, 1.7784e+00},
			{1.1872e-01, -2.0828e-02, -2.8896e-02 /*...,*/, 4.3594e-02, 7.4091e-02, 8.5981e-02},
			{4.7863e-02, 9.0866e-02, 1.3867e-01 /*...,*/, 2.2276e+00, 1.2230e+00, 1.0092e+00},
		},
		{
			{6.9574e-01, -4.4043e-01, 9.2948e-01 /*...,*/, 9.5450e-01, 7.0820e-01, 4.0560e-01},
			{2.0163e-01, 8.6028e-02, 2.1220e-01 /*...,*/, 1.5047e-01, 1.7557e-01, 1.7040e-01},
			{-4.3524e-02, 1.9791e+00, 7.8367e-01 /*...,*/, 2.0767e+00, 2.4926e+00, 2.3336e+00},
			/*...,*/
			{1.5742e+00, -3.0362e+00, -3.0963e+00 /*...,*/, 2.0745e+00, 2.6284e+00, 2.3956e+00},
			{2.2950e-01, -3.3358e-02, -9.5946e-02 /*...,*/, 1.1370e-01, 1.2174e-01, 1.1693e-01},
			{1.4589e+00, 2.4126e+00, 6.4037e-01 /*...,*/, 2.9313e+00, 2.2336e+00, 1.2770e+00},
		},
		/*...,*/
		{
			{2.5047e-01, -5.5406e-01, 1.2645e+00 /*...,*/, 6.2868e-01, -6.0493e-02, 1.5389e+00},
			{9.2031e-02, -1.4594e-01, -2.4383e-02 /*...,*/, 2.1952e-02, -6.7880e-02, 2.5287e-02},
			{1.8867e+00, 5.1775e-01, -3.5486e-01 /*...,*/, 1.2577e+00, -2.0807e+00, -1.7765e-01},
			/*...,*/
			{-1.7799e+00, -9.1624e-01, -3.7893e-01 /*...,*/, 9.4433e-01, -2.0660e+00, -1.9276e-01},
			{3.0151e-02, -7.8603e-02, 2.5738e-03 /*...,*/, 4.0625e-02, -8.5604e-02, -2.7478e-02},
			{3.2559e+00, -4.2435e-01, 2.9481e+00 /*...,*/, 3.2141e+00, -1.1062e+00, 7.0349e-01},
		},
		{
			{-2.6671e-01, -3.8844e-01, -3.0750e-01 /*...,*/, 4.5032e-01, -4.7121e-01, -1.3882e-01},
			{4.9349e-02, 1.0626e-01, 2.8606e-02 /*...,*/, -1.3075e-02, -6.4812e-02, -1.7991e-02},
			{-1.1119e+00, 1.7021e+00, 2.9340e-01 /*...,*/, 8.9383e-01, -1.1510e+00, 1.2941e-01},
			/*...,*/
			{1.8760e+00, -1.1862e+00, 1.8114e-01 /*...,*/, 8.1934e-01, -9.2910e-01, 5.6066e-02},
			{3.9879e-02, 2.4555e-02, -2.2776e-02 /*...,*/, 1.5699e-02, -3.0517e-02, -1.5159e-02},
			{2.0575e+00, 2.8044e+00, -4.2136e-01 /*...,*/, 1.1335e+00, -1.1497e+00, 6.4751e-02},
		},
		{
			{-5.6102e-01, 3.7558e-01, 2.1194e-01 /*...,*/, 4.5418e-01, -3.7228e-01, 2.6291e-01},
			{-4.2811e-02, -7.7944e-02, -1.7524e-02 /*...,*/, 6.5737e-02, -1.0265e-01, 1.1557e-01},
			{1.2875e+00, -3.9165e-01, -4.0549e-01 /*...,*/, 2.1700e+00, -3.1855e+00, 2.1169e+00},
			/*...,*/
			{-7.1299e-01, -2.3520e-01, 3.5515e-02 /*...,*/, 1.7904e+00, -2.7587e+00, 2.0914e+00},
			{-4.4265e-02, -1.5338e-02, 5.6960e-03 /*...,*/, 3.2324e-02, -5.5728e-02, 1.9095e-02},
			{3.0192e-01, 2.3755e-01, 6.2483e-01 /*...,*/, 2.3494e+00, -2.4556e+00, 1.4067e+00},
		},
	}

	if actualXq, err = actualXq.Transpose(0, 1); err != nil { // from [15, 32, 128] -> shape=[32, 15, 128] (N_Heads, sequenceLength, HeadDim)
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedXqTranspose, expectedXqTransposeSize, actualXq, 9*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	expectedKeysTransposeDims0_1_Size := []int{32, 15, 128}
	expectedKeysTransposeDims0_1 := [][][]float32{
		{
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			/*...,*/
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
		},
		{
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			/*...,*/
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
		},
		{
			{5.5868e-01, -3.3646e-01, 9.8483e-01 /*...,*/, -9.2118e-02, 2.1667e-01, 4.1570e-01},
			{1.1673e-01, 4.8191e-01, -1.4595e-01 /*...,*/, -1.1752e-01, -1.5259e-01, -1.4955e-01},
			{-6.4237e+00, 5.7450e+00, -3.4984e+00 /*...,*/, -1.4686e+00, -1.7071e+00, -1.4609e+00},
			/*...,*/
			{8.1662e+00, -1.5020e+00, -1.2152e+00 /*...,*/, -1.3395e+00, -1.5037e+00, -1.6799e+00},
			{2.0754e-01, 3.1775e-01, 1.3624e-01 /*...,*/, -7.8951e-02, -1.1078e-01, -1.2612e-01},
			{8.3640e-01, 2.8779e+00, 1.0306e+00 /*...,*/, -1.4720e+00, -1.1254e+00, -1.1091e+00},
		},
		/*...,*/
		{
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
			/*...,*/
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
		},
		{
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
			/*...,*/
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
		},
		{
			{-1.6577e-02, 5.4854e-02, -1.0922e-01 /*...,*/, 1.7037e+00, -2.0537e+00, 1.7776e+00},
			{2.5071e-01, -1.7279e-01, -8.5062e-02 /*...,*/, -8.2153e-02, 3.2857e-02, 4.5878e-02},
			{4.4648e+00, 2.7718e+00, -3.3032e-02 /*...,*/, -1.5299e+00, 1.5073e+00, -8.1492e-01},
			/*...,*/
			{-1.8412e+00, -4.8942e+00, 3.4559e+00 /*...,*/, -1.2853e+00, 1.7026e+00, -3.4815e-01},
			{1.0009e-01, -1.6104e-01, 5.1040e-03 /*...,*/, -9.7978e-02, 9.6370e-02, 1.2926e-02},
			{1.6120e-01, 6.5406e-02, 7.0973e-01 /*...,*/, 3.0597e-01, -6.9375e-01, -2.2576e+00},
		},
	}

	if actualKeys, err = actualKeys.Transpose(0, 1); err != nil { // from [5, 32, 128] -> shape=[32, 5, 128] (N_Heads, sequenceLength, HeadDim)
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedKeysTransposeDims0_1, expectedKeysTransposeDims0_1_Size, actualKeys, 12*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	expectedValuesTransposeSize := []int{32, 15, 128}
	expectedValuesTranspose := [][][]float32{
		{
			{0.0109, -0.0005, 0.0286 /*...,*/, -0.0215, -0.0026, 0.0478},
			{-0.0047, -0.0017, -0.0118 /*...,*/, -0.0085, 0.0107, -0.0003},
			{0.1022, -0.0047, 0.0568 /*...,*/, -0.0158, -0.0123, -0.0919},
			/*...,*/
			{0.0217, -0.0317, 0.0653 /*...,*/, 0.0208, 0.0455, -0.0605},
			{-0.0027, 0.0004, -0.0056 /*...,*/, -0.0052, 0.0044, -0.0110},
			{0.0008, -0.0022, 0.0154 /*...,*/, -0.0037, 0.0038, 0.1224},
		},
		{
			{0.0109, -0.0005, 0.0286 /*...,*/, -0.0215, -0.0026, 0.0478},
			{-0.0047, -0.0017, -0.0118 /*...,*/, -0.0085, 0.0107, -0.0003},
			{0.1022, -0.0047, 0.0568 /*...,*/, -0.0158, -0.0123, -0.0919},
			/*...,*/
			{0.0217, -0.0317, 0.0653 /*...,*/, 0.0208, 0.0455, -0.0605},
			{-0.0027, 0.0004, -0.0056 /*...,*/, -0.0052, 0.0044, -0.0110},
			{0.0008, -0.0022, 0.0154 /*...,*/, -0.0037, 0.0038, 0.1224},
		},
		{
			{0.0109, -0.0005, 0.0286 /*...,*/, -0.0215, -0.0026, 0.0478},
			{-0.0047, -0.0017, -0.0118 /*...,*/, -0.0085, 0.0107, -0.0003},
			{0.1022, -0.0047, 0.0568 /*...,*/, -0.0158, -0.0123, -0.0919},
			/*...,*/
			{0.0217, -0.0317, 0.0653 /*...,*/, 0.0208, 0.0455, -0.0605},
			{-0.0027, 0.0004, -0.0056 /*...,*/, -0.0052, 0.0044, -0.0110},
			{0.0008, -0.0022, 0.0154 /*...,*/, -0.0037, 0.0038, 0.1224},
		},
		/*...,*/
		{
			{0.0024, 0.0012, 0.0007 /*...,*/, 0.0007, -0.0010, 0.0004},
			{-0.0002, 0.0029, -0.0137 /*...,*/, 0.0064, -0.0113, 0.0087},
			{-0.0246, -0.0609, 0.0164 /*...,*/, 0.0708, 0.0168, -0.0362},
			/*...,*/
			{-0.0386, -0.0109, -0.0181 /*...,*/, 0.0031, -0.0533, 0.0256},
			{-0.0118, 0.0015, 0.0078 /*...,*/, 0.0035, -0.0044, 0.0066},
			{-0.0097, -0.0115, 0.0159 /*...,*/, 0.0167, -0.0084, -0.0090},
		},

		{
			{0.0024, 0.0012, 0.0007 /*...,*/, 0.0007, -0.0010, 0.0004},
			{-0.0002, 0.0029, -0.0137 /*...,*/, 0.0064, -0.0113, 0.0087},
			{-0.0246, -0.0609, 0.0164 /*...,*/, 0.0708, 0.0168, -0.0362},
			/*...,*/
			{-0.0386, -0.0109, -0.0181 /*...,*/, 0.0031, -0.0533, 0.0256},
			{-0.0118, 0.0015, 0.0078 /*...,*/, 0.0035, -0.0044, 0.0066},
			{-0.0097, -0.0115, 0.0159 /*...,*/, 0.0167, -0.0084, -0.0090},
		},
		{
			{0.0024, 0.0012, 0.0007 /*...,*/, 0.0007, -0.0010, 0.0004},
			{-0.0002, 0.0029, -0.0137 /*...,*/, 0.0064, -0.0113, 0.0087},
			{-0.0246, -0.0609, 0.0164 /*...,*/, 0.0708, 0.0168, -0.0362},
			/*...,*/
			{-0.0386, -0.0109, -0.0181 /*...,*/, 0.0031, -0.0533, 0.0256},
			{-0.0118, 0.0015, 0.0078 /*...,*/, 0.0035, -0.0044, 0.0066},
			{-0.0097, -0.0115, 0.0159 /*...,*/, 0.0167, -0.0084, -0.0090},
		},
	}

	if actualValues, err = actualValues.Transpose(0, 1); err != nil { // from [15, 32, 128] -> shape=[32, 15, 128] (N_Heads, sequenceLength, HeadDim)
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedValuesTranspose, expectedValuesTransposeSize, actualValues, common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	expectedKeysTransposeDims1_2_Size := []int{32, 128, 15}
	expectedKeysTransposeDims1_2 := [][][]float32{
		{
			{5.5868e-01, 1.1673e-01, -6.4237e+00 /*...,*/, 8.1662e+00, 2.0754e-01, 8.3640e-01},
			{-3.3646e-01, 4.8191e-01, 5.7450e+00 /*...,*/, -1.5020e+00, 3.1775e-01, 2.8779e+00},
			{9.8483e-01, -1.4595e-01, -3.4984e+00 /*...,*/, -1.2152e+00, 1.3624e-01, 1.0306e+00},
			/*...,*/
			{-9.2118e-02, -1.1752e-01, -1.4686e+00 /*...,*/, -1.3395e+00, -7.8951e-02, -1.4720e+00},
			{2.1667e-01, -1.5259e-01, -1.7071e+00 /*...,*/, -1.5037e+00, -1.1078e-01, -1.1254e+00},
			{4.1570e-01, -1.4955e-01, -1.4609e+00 /*...,*/, -1.6799e+00, -1.2612e-01, -1.1091e+00},
		},
		{
			{5.5868e-01, 1.1673e-01, -6.4237e+00 /*...,*/, 8.1662e+00, 2.0754e-01, 8.3640e-01},
			{-3.3646e-01, 4.8191e-01, 5.7450e+00 /*...,*/, -1.5020e+00, 3.1775e-01, 2.8779e+00},
			{9.8483e-01, -1.4595e-01, -3.4984e+00 /*...,*/, -1.2152e+00, 1.3624e-01, 1.0306e+00},
			/*...,*/
			{-9.2118e-02, -1.1752e-01, -1.4686e+00 /*...,*/, -1.3395e+00, -7.8951e-02, -1.4720e+00},
			{2.1667e-01, -1.5259e-01, -1.7071e+00 /*...,*/, -1.5037e+00, -1.1078e-01, -1.1254e+00},
			{4.1570e-01, -1.4955e-01, -1.4609e+00 /*...,*/, -1.6799e+00, -1.2612e-01, -1.1091e+00},
		},
		{
			{5.5868e-01, 1.1673e-01, -6.4237e+00 /*...,*/, 8.1662e+00, 2.0754e-01, 8.3640e-01},
			{-3.3646e-01, 4.8191e-01, 5.7450e+00 /*...,*/, -1.5020e+00, 3.1775e-01, 2.8779e+00},
			{9.8483e-01, -1.4595e-01, -3.4984e+00 /*...,*/, -1.2152e+00, 1.3624e-01, 1.0306e+00},
			/*...,*/
			{-9.2118e-02, -1.1752e-01, -1.4686e+00 /*...,*/, -1.3395e+00, -7.8951e-02, -1.4720e+00},
			{2.1667e-01, -1.5259e-01, -1.7071e+00 /*...,*/, -1.5037e+00, -1.1078e-01, -1.1254e+00},
			{4.1570e-01, -1.4955e-01, -1.4609e+00 /*...,*/, -1.6799e+00, -1.2612e-01, -1.1091e+00},
		},
		/*...,*/
		{
			{-1.6577e-02, 2.5071e-01, 4.4648e+00 /*...,*/, -1.8412e+00, 1.0009e-01, 1.6120e-01},
			{5.4854e-02, -1.7279e-01, 2.7718e+00 /*...,*/, -4.8942e+00, -1.6104e-01, 6.5406e-02},
			{-1.0922e-01, -8.5062e-02, -3.3032e-02 /*...,*/, 3.4559e+00, 5.1040e-03, 7.0973e-01},
			/*...,*/
			{1.7037e+00, -8.2153e-02, -1.5299e+00 /*...,*/, -1.2853e+00, -9.7978e-02, 3.0597e-01},
			{-2.0537e+00, 3.2857e-02, 1.5073e+00 /*...,*/, 1.7026e+00, 9.6370e-02, -6.9375e-01},
			{1.7776e+00, 4.5878e-02, -8.1492e-01 /*...,*/, -3.4815e-01, 1.2926e-02, -2.2576e+00},
		},
		{
			{-1.6577e-02, 2.5071e-01, 4.4648e+00 /*...,*/, -1.8412e+00, 1.0009e-01, 1.6120e-01},
			{5.4854e-02, -1.7279e-01, 2.7718e+00 /*...,*/, -4.8942e+00, -1.6104e-01, 6.5406e-02},
			{-1.0922e-01, -8.5062e-02, -3.3032e-02 /*...,*/, 3.4559e+00, 5.1040e-03, 7.0973e-01},
			/*...,*/
			{1.7037e+00, -8.2153e-02, -1.5299e+00 /*...,*/, -1.2853e+00, -9.7978e-02, 3.0597e-01},
			{-2.0537e+00, 3.2857e-02, 1.5073e+00 /*...,*/, 1.7026e+00, 9.6370e-02, -6.9375e-01},
			{1.7776e+00, 4.5878e-02, -8.1492e-01 /*...,*/, -3.4815e-01, 1.2926e-02, -2.2576e+00},
		},
		{
			{-1.6577e-02, 2.5071e-01, 4.4648e+00 /*...,*/, -1.8412e+00, 1.0009e-01, 1.6120e-01},
			{5.4854e-02, -1.7279e-01, 2.7718e+00 /*...,*/, -4.8942e+00, -1.6104e-01, 6.5406e-02},
			{-1.0922e-01, -8.5062e-02, -3.3032e-02 /*...,*/, 3.4559e+00, 5.1040e-03, 7.0973e-01},
			/*...,*/
			{1.7037e+00, -8.2153e-02, -1.5299e+00 /*...,*/, -1.2853e+00, -9.7978e-02, 3.0597e-01},
			{-2.0537e+00, 3.2857e-02, 1.5073e+00 /*...,*/, 1.7026e+00, 9.6370e-02, -6.9375e-01},
			{1.7776e+00, 4.5878e-02, -8.1492e-01 /*...,*/, -3.4815e-01, 1.2926e-02, -2.2576e+00},
		},
	}

	if actualKeys, err = actualKeys.Transpose(1, 2); err != nil { // from [32, 15, 128] -> shape=[32, 128, 15] (N_Heads, HeadDim, sequenceLength)
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedKeysTransposeDims1_2, expectedKeysTransposeDims1_2_Size, actualKeys, 12*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		Goal in Python manner:
		scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
	*/

	expectedScoresSize := []int{32, 15, 15}
	expectedScores := [][][]float32{
		{
			{3.1636e+00, 4.1748e-02, 7.1905e-01 /*...,*/, 6.7906e-01, -1.0255e-02, 1.6342e+00},
			{5.3825e-01, 1.7084e-02, 2.2871e-01 /*...,*/, 1.8474e-01, 3.5051e-03, 2.5659e-01},
			{8.2943e+00, 2.8990e-01, 6.1792e+00 /*...,*/, 4.6819e+00, 8.8838e-02, 5.3201e+00},
			/*...,*/
			{6.6150e+00, 1.5766e-01, 4.4712e+00 /*...,*/, 6.0474e+00, 6.7458e-02, 5.4719e+00},
			{2.5776e-01, 9.1109e-03, 1.9566e-01 /*...,*/, 3.7384e-01, 1.1165e-02, 2.9912e-01},
			{4.6241e+00, 1.3076e-01, 3.8184e+00 /*...,*/, 6.0326e+00, 1.0872e-01, 6.0652e+00},
		},
		{
			{2.7323e+00, 9.0661e-02, 1.3101e+00 /*...,*/, 1.2840e+00, 1.2566e-02, 1.4027e+00},
			{5.0203e-01, 7.2063e-03, 1.9911e-01 /*...,*/, 2.1079e-01, 4.2153e-03, 3.1311e-01},
			{7.2577e+00, 1.0232e-01, 3.1345e+00 /*...,*/, 2.6152e+00, 1.7088e-02, 4.3419e+00},
			/*...,*/
			{7.1406e+00, 1.5857e-02, 1.6561e+00 /*...,*/, 3.1646e+00, 7.0212e-02, 5.0311e+00},
			{2.8141e-01, -4.7072e-04, -3.5377e-03 /*...,*/, 2.7828e-01, 4.7701e-03, 2.5944e-01},
			{7.3885e+00, 1.7566e-01, 4.0247e+00 /*...,*/, 4.9096e+00, 1.5687e-01, 5.4150e+00},
		},
		{
			{2.6903e+00, 2.0476e-02, -1.8658e-02 /*...,*/, 6.1869e-01, -2.6446e-02, 7.9876e-01},
			{4.4384e-01, 1.1797e-02, 1.2359e-01 /*...,*/, 5.8250e-02, -8.9389e-03, 6.9690e-02},
			{6.0182e+00, 2.1051e-01, 4.9419e+00 /*...,*/, -8.2138e-01, -1.6082e-01, 1.1748e+00},
			/*...,*/
			{3.9564e+00, -3.1563e-01, -2.4268e+00 /*...,*/, 6.1895e+00, -1.7940e-01, 3.7642e+00},
			{1.7185e-01, -6.4635e-03, -8.6048e-02 /*...,*/, 6.0917e-01, 1.2682e-02, 4.3847e-01},
			{2.9314e+00, -2.1315e-01, -1.6156e+00 /*...,*/, 3.6495e+00, 1.4751e-01, 5.5472e+00},
		},
		/*...,*/
		{
			{2.0774e+00, 1.2698e-01, 6.9547e-01 /*...,*/, 5.9082e-01, 9.7963e-03, -2.0268e+00},
			{2.4089e-01, 7.4250e-03, -2.2346e-02 /*...,*/, -5.6294e-02, -3.2739e-03, 9.2240e-02},
			{8.3201e+00, -1.1158e-03, -1.7454e+00 /*...,*/, -4.9652e+00, -2.7248e-01, 2.6623e+00},
			/*...,*/
			{7.8913e+00, -1.2278e-01, -5.0430e+00 /*...,*/, -1.5126e+00, -1.9943e-01, 1.8527e+00},
			{3.5490e-01, -9.5775e-03, -3.1874e-01 /*...,*/, -8.0427e-02, -9.8166e-03, 1.7120e-01},
			{8.7407e+00, -1.9788e-01, -2.8999e+00 /*...,*/, -7.7589e-01, -1.1322e-01, -2.2040e+00},
		},
		{
			{3.0519e+00, -9.1138e-02, -2.8181e+00 /*...,*/, -1.5736e+00, -1.0311e-01, 6.1995e-01},
			{1.7599e-01, 8.5416e-03, 1.4766e-01 /*...,*/, -1.4935e-01, -2.9988e-03, 7.5232e-02},
			{5.7762e+00, 2.7664e-02, 1.3447e+00 /*...,*/, -2.7985e+00, -1.2605e-01, 3.6684e-01},
			/*...,*/
			{5.2272e+00, -3.4942e-02, -2.7111e+00 /*...,*/, 1.9022e+00, 1.1050e-01, 1.4746e+00},
			{1.1987e-01, -2.3555e-03, -6.9269e-02 /*...,*/, -6.1157e-03, 1.9582e-03, 6.3316e-02},
			{6.5878e+00, -1.2186e-01, -2.9440e+00 /*...,*/, -3.6560e+00, -1.2001e-01, 3.8644e+00},
		},
		{
			{4.2182e+00, -1.8271e-01, -4.3795e+00 /*...,*/, -2.5793e+00, -1.7767e-01, 8.9631e-01},
			{6.2380e-01, -1.8290e-02, -5.6274e-01 /*...,*/, -3.5090e-01, -3.0627e-02, 1.0796e-02},
			{1.8260e+01, -6.0239e-01, -1.4060e+01 /*...,*/, -1.2341e+01, -9.8128e-01, 1.7132e+00},
			/*...,*/
			{1.6575e+01, -6.2350e-01, -1.4134e+01 /*...,*/, -1.0294e+01, -8.8164e-01, 1.5936e+00},
			{3.0378e-01, -1.1710e-02, -2.6946e-01 /*...,*/, -1.7485e-01, -1.6378e-02, -1.4945e-02},
			{2.1142e+01, -9.2544e-01, -1.9238e+01 /*...,*/, -1.5197e+01, -1.1620e+00, 6.3661e+00},
		},
	}

	xqMatMulKeys, err := ml.MatMul(actualXq, actualKeys) // matmul([32,15,128], [32,128,15]) -> shape=[32,15,15] (N_Heads, sequenceLength, sequenceLength)
	if err != nil {
		t.Fatal(err)
	}
	actualScores, err := ml.DivToScalar(xqMatMulKeys, dtype.BFloat16fromFloat32(float32(math.Sqrt(float64(attention.HeadDim))))) // shape=[32,15,15]
	if err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedScores, expectedScoresSize, actualScores, 70*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}
	negInf := float32(math.Inf(-1))

	if mask != nil {
		expectedScoresPlusMaskSize := []int{32, 15, 15}
		expectedScoresPlusMask := [][][]float32{
			{
				{3.1636e+00, negInf, negInf /*...,*/, negInf, negInf, negInf},
				{5.3825e-01, 1.7084e-02, negInf /*...,*/, negInf, negInf, negInf},
				{8.2943e+00, 2.8990e-01, 6.1792e+00 /*...,*/, negInf, negInf, negInf},
				/*...,*/
				{6.6150e+00, 1.5766e-01, 4.4712e+00 /*...,*/, 6.0474e+00, negInf, negInf},
				{2.5776e-01, 9.1109e-03, 1.9566e-01 /*...,*/, 3.7384e-01, 1.1165e-02, negInf},
				{4.6241e+00, 1.3076e-01, 3.8184e+00 /*...,*/, 6.0326e+00, 1.0872e-01, 6.0652e+00},
			},
			{
				{2.7323e+00, negInf, negInf /*...,*/, negInf, negInf, negInf},
				{5.0203e-01, 7.2063e-03, negInf /*...,*/, negInf, negInf, negInf},
				{7.2577e+00, 1.0232e-01, 3.1345e+00 /*...,*/, negInf, negInf, negInf},
				/*...,*/
				{7.1406e+00, 1.5857e-02, 1.6561e+00 /*...,*/, 3.1646e+00, negInf, negInf},
				{2.8141e-01, -4.7072e-04, -3.5377e-03 /*...,*/, 2.7828e-01, 4.7701e-03, negInf},
				{7.3885e+00, 1.7566e-01, 4.0247e+00 /*...,*/, 4.9096e+00, 1.5687e-01, 5.4150e+00},
			},
			{
				{2.6903e+00, negInf, negInf /*...,*/, negInf, negInf, negInf},
				{4.4384e-01, 1.1797e-02, negInf /*...,*/, negInf, negInf, negInf},
				{6.0182e+00, 2.1051e-01, 4.9419e+00 /*...,*/, negInf, negInf, negInf},
				/*...,*/
				{3.9564e+00, -3.1563e-01, -2.4268e+00 /*...,*/, 6.1895e+00, negInf, negInf},
				{1.7185e-01, -6.4635e-03, -8.6048e-02 /*...,*/, 6.0917e-01, 1.2682e-02, negInf},
				{2.9314e+00, -2.1315e-01, -1.6156e+00 /*...,*/, 3.6495e+00, 1.4751e-01, 5.5472e+00},
			},
			/*...,*/
			{
				{2.0774e+00, negInf, negInf /*...,*/, negInf, negInf, negInf},
				{2.4089e-01, 7.4250e-03, negInf /*...,*/, negInf, negInf, negInf},
				{8.3201e+00, -1.1158e-03, -1.7454e+00 /*...,*/, negInf, negInf, negInf},
				/*...,*/
				{7.8913e+00, -1.2278e-01, -5.0430e+00 /*...,*/, -1.5126e+00, negInf, negInf},
				{3.5490e-01, -9.5775e-03, -3.1874e-01 /*...,*/, -8.0427e-02, -9.8166e-03, negInf},
				{8.7407e+00, -1.9788e-01, -2.8999e+00 /*...,*/, -7.7589e-01, -1.1322e-01, -2.2040e+00},
			},
			{
				{3.0519e+00, negInf, negInf /*...,*/, negInf, negInf, negInf},
				{1.7599e-01, 8.5416e-03, negInf /*...,*/, negInf, negInf, negInf},
				{5.7762e+00, 2.7664e-02, 1.3447e+00 /*...,*/, negInf, negInf, negInf},
				/*...,*/
				{5.2272e+00, -3.4942e-02, -2.7111e+00 /*...,*/, 1.9022e+00, negInf, negInf},
				{1.1987e-01, -2.3555e-03, -6.9269e-02 /*...,*/, -6.1157e-03, 1.9582e-03, negInf},
				{6.5878e+00, -1.2186e-01, -2.9440e+00 /*...,*/, -3.6560e+00, -1.2001e-01, 3.8644e+00},
			},
			{
				{4.2182e+00, negInf, negInf /*...,*/, negInf, negInf, negInf},
				{6.2380e-01, -1.8290e-02, negInf /*...,*/, negInf, negInf, negInf},
				{1.8260e+01, -6.0239e-01, -1.4060e+01 /*...,*/, negInf, negInf, negInf},
				/*...,*/
				{1.6575e+01, -6.2350e-01, -1.4134e+01 /*...,*/, -1.0294e+01, negInf, negInf},
				{3.0378e-01, -1.1710e-02, -2.6946e-01 /*...,*/, -1.7485e-01, -1.6378e-02, negInf},
				{2.1142e+01, -9.2544e-01, -1.9238e+01 /*...,*/, -1.5197e+01, -1.1620e+00, 6.3661e+00},
			},
		}
		if actualScores, err = ml.Add(actualScores, mask); err != nil { // shape=[32,5,5]
			t.Fatal(err)
		}
		if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedScoresPlusMask, expectedScoresPlusMaskSize, actualScores, 70*common.THRESHOLD_BF16, true); err != nil {
			t.Fatal(err)
		}
	}

	/*
		Goal in Python manner:
		scores = F.softmax(scores.float(), dim=-1).type_as(xq)
	*/

	expectedScoresSoftmaxSize := []int{32, 15, 15}
	expectedScoresSoftmax := [][][]float32{
		{
			{1.0000e+00, 0.0000e+00, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{6.2742e-01, 3.7258e-01, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{8.9210e-01, 2.9794e-04, 1.0760e-01 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			/*...,*/
			{2.1860e-01, 3.4297e-04, 2.5621e-02 /*...,*/, 1.2392e-01, 0.0000e+00, 0.0000e+00},
			{7.8082e-02, 6.0892e-02, 7.3380e-02 /*...,*/, 8.7692e-02, 6.1017e-02, 0.0000e+00},
			{5.6210e-02, 6.2858e-04, 2.5111e-02 /*...,*/, 2.2987e-01, 6.1488e-04, 2.3750e-01},
		},
		{
			{1.0000e+00, 0.0000e+00, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{6.2124e-01, 3.7876e-01, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{9.8331e-01, 7.6762e-04, 1.5923e-02 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			/*...,*/
			{7.9297e-01, 6.3827e-04, 3.2911e-03 /*...,*/, 1.4875e-02, 0.0000e+00, 0.0000e+00},
			{8.5278e-02, 6.4330e-02, 6.4133e-02 /*...,*/, 8.5012e-02, 6.4668e-02, 0.0000e+00},
			{5.0658e-01, 3.7337e-04, 1.7528e-02 /*...,*/, 4.2469e-02, 3.6642e-04, 7.0400e-02},
		},
		{
			{1.0000e+00, 0.0000e+00, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{6.0636e-01, 3.9364e-01, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{7.4413e-01, 2.2357e-03, 2.5363e-01 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			/*...,*/
			{8.9983e-02, 1.2556e-03, 1.5205e-04 /*...,*/, 8.3941e-01, 0.0000e+00, 0.0000e+00},
			{7.5333e-02, 6.3029e-02, 5.8207e-02 /*...,*/, 1.1666e-01, 6.4247e-02, 0.0000e+00},
			{5.8559e-02, 2.5229e-03, 6.2063e-04 /*...,*/, 1.2007e-01, 3.6186e-03, 8.0090e-01},
		},
		/*...,*/
		{
			{1.0000e+00, 0.0000e+00, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{5.5810e-01, 4.4190e-01, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{9.9971e-01, 2.4323e-04, 4.2508e-05 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			/*...,*/
			{9.7594e-01, 3.2282e-04, 2.3557e-06 /*...,*/, 8.0421e-05, 0.0000e+00, 0.0000e+00},
			{1.0354e-01, 7.1918e-02, 5.2792e-02 /*...,*/, 6.6999e-02, 7.1901e-02, 0.0000e+00},
			{9.9115e-01, 1.3007e-04, 8.7239e-06 /*...,*/, 7.2972e-05, 1.4156e-04, 1.7496e-05},
		},
		{
			{1.0000e+00, 0.0000e+00, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{5.4177e-01, 4.5823e-01, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{9.8514e-01, 3.1402e-03, 1.1720e-02 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			/*...,*/
			{8.7142e-01, 4.5174e-03, 3.1091e-04 /*...,*/, 3.1344e-02, 0.0000e+00, 0.0000e+00},
			{7.8953e-02, 6.9870e-02, 6.5348e-02 /*...,*/, 6.9608e-02, 7.0172e-02, 0.0000e+00},
			{9.1381e-01, 1.1140e-03, 6.6258e-05 /*...,*/, 3.2509e-05, 1.1160e-03, 5.9992e-02},
		},
		{
			{1.0000e+00, 0.0000e+00, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{6.5523e-01, 3.4477e-01, 0.0000e+00 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			{1.0000e+00, 6.4317e-09, 9.1998e-15 /*...,*/, 0.0000e+00, 0.0000e+00, 0.0000e+00},
			/*...,*/
			{1.0000e+00, 3.3958e-08, 4.6058e-14 /*...,*/, 2.1435e-12, 0.0000e+00, 0.0000e+00},
			{9.9469e-02, 7.2556e-02, 5.6071e-02 /*...,*/, 6.1635e-02, 7.2218e-02, 0.0000e+00},
			{1.0000e+00, 2.6086e-10, 2.9067e-18 /*...,*/, 1.6528e-16, 2.0591e-10, 3.8288e-07},
		},
	}

	actualScores32, err := actualScores.ToFloat32() // shape=[32,15,15] dtype=DT_F32
	if err != nil {
		t.Fatal(err)
	}
	if actualScores32, err = ml.Softmax(actualScores32, len(actualScores32.Size)-1); err != nil { // shape=[32,15,15] dtype=DT_F32
		t.Fatal(err)
	}
	if actualScores, err = actualScores32.ToBFloat16(); err != nil { // shape=[32,15,15] (N_Heads, sequenceLength, sequenceLength) dtype=DT_BF16
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedScoresSoftmax, expectedScoresSoftmaxSize, actualScores, 2*common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		Goal in Python manner:
		output = torch.matmul(scores, values)
		output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
	*/

	expectedOutputBeforeWeightsSize := []int{15, 4096}
	expectedOutputBeforeWeights := [][]float32{
		{1.0913e-02, -4.7801e-04, 2.8640e-02 /*...,*/, 6.6461e-04, -9.7815e-04, 3.5757e-04},
		{5.0847e-03, -9.3381e-04, 1.3565e-02 /*...,*/, 2.6314e-03, -4.5329e-03, 3.2234e-03},
		{2.0728e-02, -9.3684e-04, 3.1660e-02 /*...,*/, 6.6461e-04, -9.7815e-04, 3.5757e-04},
		/*...,*/
		{8.2263e-03, -1.8792e-03, 3.4125e-02 /*...,*/, 6.6462e-04, -9.7816e-04, 3.5757e-04},
		{1.8169e-02, -5.5701e-05, 1.4084e-02 /*...,*/, -3.6742e-04, -7.6021e-03, -8.4274e-04},
		{1.2438e-02, -5.4218e-03, 3.1450e-02 /*...,*/, 6.6462e-04, -9.7816e-04, 3.5757e-04},
	}

	actualOutput, err := ml.MatMul(actualScores, actualValues)
	if err != nil {
		t.Fatal(err)
	}
	if actualOutput, err = actualOutput.Transpose(0, 1); err != nil {
		t.Fatal(err)
	}
	outputTrailingSize := actualOutput.GetElementCount() / sequenceLength
	if actualOutput, err = actualOutput.Reshape([]int{sequenceLength, outputTrailingSize}); err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedOutputBeforeWeights, expectedOutputBeforeWeightsSize, actualOutput, common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	expectedOutputAfterWeightsSize := []int{15, 4096}
	expectedOutputAfterWeights := [][]float32{
		{2.1353e-03, -1.7818e-03, -1.6321e-04 /*...,*/, 2.6612e-03, -4.9833e-04, -1.1525e-03},
		{1.0629e-03, -1.8515e-03, 1.0656e-04 /*...,*/, 5.2467e-04, -3.2402e-04, -3.0441e-05},
		{1.7809e-03, 3.5385e-04, 7.8608e-04 /*...,*/, 7.0986e-03, 2.6803e-03, -3.2171e-03},
		/*...,*/
		{3.4250e-03, 1.5080e-03, 4.7115e-03 /*...,*/, 3.1859e-03, 1.5443e-04, -8.0089e-04},
		{2.9352e-03, -2.7936e-03, 5.4721e-03 /*...,*/, 1.4995e-03, -2.1091e-03, -1.0560e-03},
		{7.1176e-03, -2.5008e-03, -1.6193e-03 /*...,*/, -1.0986e-02, -3.1216e-03, -6.2308e-03},
	}
	/*
		Apply lat.attn_wo weights to output
	*/

	// lat.attn_wo: [out_features, in_features] -> shape: [4096 4096] -> [N_Heads * HeadDim, Dim]
	if actualOutput, err = ml.LinearTransformation(actualOutput, attention.attn_wo); err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedOutputAfterWeights, expectedOutputAfterWeightsSize, actualOutput, common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	return actualOutput
}

func testTransformerBlock_FeedForward_Forward(t *testing.T, feedForward *LlamaFeedForward, x *ml.Tensor) *ml.Tensor {
	/*
		Goal in Python manner:
		self.w2(F.silu(self.w1(x)) * self.w3(x))
		-->
		self.ffn_down(F.silu(self.ffn_gate(x)) * self.ffn_up(x))
	*/
	h, err := ml.LinearTransformation(x, feedForward.ffn_gate)
	if err != nil {
		t.Fatal(err)
	}
	if h, err = ml.Silu(h); err != nil {
		t.Fatal(err)
	}
	ffnUpX, err := ml.LinearTransformation(x, feedForward.ffn_up)
	if err != nil {
		t.Fatal(err)
	}
	if h, err = ml.MultiplyElementwise(h, ffnUpX); err != nil {
		t.Fatal(err)
	}
	actualOutput, err := ml.LinearTransformation(h, feedForward.ffn_down)
	if err != nil {
		t.Fatal(err)
	}
	return actualOutput
}

func testTransformerBlock_Forward(t *testing.T, skipCompareTestTensor bool, infContext *InferenceContext, transformerBlock *LlamaTransformerBlock, x *ml.Tensor, startPos int, freqsCis *ml.Tensor, mask *ml.Tensor) *ml.Tensor {
	/*
		h, err := ltb.attention.Forward(infContext, normalizedX, startPos, freqsCis, mask)
	*/
	normalizedX := testTransformerBlock_AttnNorm_Forward(t, skipCompareTestTensor, transformerBlock, x)
	h := testTransformerBlock_Attention_Forward(t, skipCompareTestTensor, infContext, transformerBlock.attention, normalizedX, startPos, freqsCis, mask)

	expectedHBeforeFeedForwardSize := []int{15, 4096}
	expectedHBeforeFeedForward := [][]float32{
		{2.4004e-03, -2.2815e-03, -7.4686e-04 /*...,*/, 6.4759e-03, -4.3491e-04, 3.7711e-05},
		{8.9789e-04, -2.0947e-03, 2.7059e-04 /*...,*/, 3.7304e-04, 2.6935e-05, 7.0198e-04},
		{5.2905e-03, 7.5560e-03, 8.3949e-04 /*...,*/, 6.3738e-03, -7.9398e-03, -2.3894e-03},
		/*...,*/
		{-6.3406e-03, -1.9558e-03, 6.5730e-03 /*...,*/, -3.9552e-03, -4.1486e-03, 7.8051e-03},
		{2.4736e-03, -3.1865e-03, 5.4655e-03 /*...,*/, 8.7388e-04, -2.6126e-03, -3.8843e-04},
		{4.2490e-03, 1.3901e-03, -1.7929e-03 /*...,*/, -1.0178e-02, -2.6181e-03, -3.9267e-03},
	}

	var err error
	if h, err = ml.Add(x, h); err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedHBeforeFeedForward, expectedHBeforeFeedForwardSize, h, common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}

	/*
		h, err = ltb.ffn_norm.Forward(context, h)
	*/
	expectedOutputSize := []int{15, 4096}
	expectedOutput := [][]float32{
		{2.2415e-03, 4.0932e-03, -6.6109e-04 /*...,*/, 1.9012e-02, -4.2205e-03, -2.5585e-03},
		{3.1618e-03, 2.3711e-02, 1.2093e-02 /*...,*/, 9.0826e-03, -5.2393e-02, -5.0278e-03},
		{1.7812e-02, 1.6469e-02, 2.6643e-02 /*...,*/, -4.1046e-03, 2.1696e-03, -1.1988e-02},
		/*...,*/
		{-2.7957e-02, -2.1291e-02, 6.6948e-03 /*...,*/, -2.3827e-02, 1.2252e-02, -6.3764e-03},
		{8.9129e-03, 4.0429e-03, -7.4083e-03 /*...,*/, -1.5316e-03, -5.1894e-03, -9.4591e-03},
		{1.0622e-02, 1.0772e-02, 7.2363e-03 /*...,*/, -1.5355e-02, 1.5474e-02, -9.4875e-05},
	}
	normalizedH, err := transformerBlock.ffn_norm.Forward(infContext, h)
	if err != nil {
		t.Fatal(err)
	}
	ffnOutput := testTransformerBlock_FeedForward_Forward(t, transformerBlock.feedForward, normalizedH)
	actualOutput, err := ml.Add(h, ffnOutput)
	if err != nil {
		t.Fatal(err)
	}
	if err := ml.CompareTestTensorSkippable(skipCompareTestTensor, expectedOutput, expectedOutputSize, actualOutput, common.THRESHOLD_BF16, true); err != nil {
		t.Fatal(err)
	}
	return actualOutput
}

func testTransformer_Forward(t *testing.T, onlyFirstLayer bool, infContext *InferenceContext, transformer *LlamaTransformer, inputTokens *ml.Tensor, startPos int) *ml.Tensor {
	skipCompareTestTensor := !onlyFirstLayer || startPos > 0
	var err error
	actualInputTensor, actualFreqsCis, actualMask := testTransformer_Prepare(t, skipCompareTestTensor, transformer, inputTokens, startPos)

	currentTensor := actualInputTensor
	if onlyFirstLayer {
		firstLayer := transformer.Layers[0]
		currentTensor = testTransformerBlock_Forward(t, skipCompareTestTensor, infContext, firstLayer, currentTensor, startPos, actualFreqsCis, actualMask)
	} else {
		for layerIdx, layer := range transformer.Layers {
			infContext.Logf("Running transformer block layer: %d / %d\n", layerIdx+1, len(transformer.Layers))
			currentTensor = testTransformerBlock_Forward(t, true, infContext, layer, currentTensor, startPos, actualFreqsCis, actualMask)
		}
	}
	if currentTensor, err = transformer.output_norm.Forward(infContext, currentTensor); err != nil {
		t.Fatal(err)
	}
	output, err := ml.LinearTransformation(currentTensor, transformer.output)
	if err != nil {
		t.Fatal(err)
	}
	if output, err = output.ToFloat32(); err != nil {
		t.Fatal(err)
	}
	return output
}

func testSimulatedLog(format string, v ...any) {
	fmt.Printf(format, v...)
}

func testSimulatedInternal(t *testing.T, onlyFirstLayer bool) {
	modelDir := "../../models-original/Meta-Llama-3.1-8B-Instruct"
	if _, err := os.Stat(modelDir); err != nil {
		t.Skipf("Model directory \"%s\" is not found, passing this test: %s", modelDir, "TestSimulated")
		return
	}
	var err error
	common.GLogger, err = common.NewLogger(os.Stdout, nil)
	if err != nil {
		panic(err)
	}
	defer common.GLogger.Close()

	llamaModel, err := LoadModel(modelDir)
	if err != nil {
		t.Fatal(err)
	}
	defer llamaModel.Free()

	/*
	   promptTokens: " <|begin_of_text|><|start_header_id|>user<|end_header_id|>

	   What is your name?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

	   "
	*/
	promptTokens := []TokenId{128000, 128006, 882, 128007, 271, 3923, 374, 701, 836, 30, 128009, 128006, 78191, 128007, 271}

	inferenceArgs := common.NewInferenceArgs()
	inferenceArgs.SequenceLength = 20
	infContext := NewInferenceContext(llamaModel, inferenceArgs, testSimulatedLog)

	tokens, err := ml.Full([]int{infContext.SequenceLength}, ml.DT_INT32, int32(llamaModel.Vocabulary.PadId))
	if err != nil {
		t.Fatal(err)
	}
	for i, token := range promptTokens {
		if err := tokens.SetItem([]int{i}, int32(token)); err != nil {
			t.Fatal(err)
		}
	}

	prevPos := 0
	minPromptLength := len(promptTokens)
	isFirstIteration := true
	for curPos := minPromptLength; curPos < infContext.SequenceLength; curPos++ {
		inputTokensSlice, err := tokens.Slice([]int{prevPos}, []int{curPos})
		if err != nil {
			t.Fatal(err)
		}

		actualLogits := testTransformer_Forward(t, onlyFirstLayer, infContext, llamaModel.Transformer, inputTokensSlice, prevPos)
		if onlyFirstLayer && isFirstIteration {
			// Although it is a valid and significant output as a tensor, it isn't a meaningful outcome,
			// because while producing this output, only first attention layer was run, not completely 32 of them.
			// But it is valuable to validate the whole model flow mathematically,
			// even if we had some precision differences between ours and original Llama Python code.

			expectedLogitsOnlyFirstLayerSize := []int{15, 128256}
			expectedLogitsOnlyFirstLayer := [][]float32{
				{-2.5667, -3.9571, -2.3716 /*...,*/, 1.8254, 1.8255, 1.8255},
				{0.9026, -1.0940, -1.1630 /*...,*/, 0.0943, 0.0941, 0.0939},
				{-0.4689, 0.9260, -0.0103 /*...,*/, -0.0544, -0.0547, -0.0547},
				/*...,*/
				{3.4908, 2.4273, 2.3383 /*...,*/, -1.1902, -1.1906, -1.1905},
				{0.2826, -0.5322, -1.1386 /*...,*/, -0.4629, -0.4632, -0.4632},
				{1.2469, 1.5909, 0.0247 /*...,*/, -0.5020, -0.5021, -0.5022},
			}
			if err := ml.CompareTestTensorSkippable(!onlyFirstLayer, expectedLogitsOnlyFirstLayer, expectedLogitsOnlyFirstLayerSize, actualLogits, 30*common.THRESHOLD_BF16, true); err != nil {
				t.Fatal(err)
			}
		}
		if actualLogits, err = actualLogits.Slice([]int{actualLogits.Size[0] - 1}, []int{actualLogits.Size[0]}); err != nil {
			t.Fatal(err)
		}
		if onlyFirstLayer && isFirstIteration {
			expectedLogitsLastRowOnlyFirstLayerSize := []int{1, 128256}
			expectedLogitsLastRowOnlyFirstLayer := [][]float32{
				{1.2469, 1.5909, 0.0247 /*...,*/, -0.5020, -0.5021, -0.5022},
			}
			if err := ml.CompareTestTensorSkippable(!onlyFirstLayer, expectedLogitsLastRowOnlyFirstLayer, expectedLogitsLastRowOnlyFirstLayerSize, actualLogits, 30*common.THRESHOLD_BF16, true); err != nil {
				t.Fatal(err)
			}
		}
		actualNextToken, err := ml.Argmax(actualLogits, len(actualLogits.Size)-1) // shape=[1,1] dtype=DT_INT32
		if err != nil {
			t.Fatal(err)
		}
		if onlyFirstLayer && isFirstIteration {
			expectedNextTokenOnlyFirstLayerSize := []int{1}
			expectedNextTokenOnlyFirstLayer := []float32{
				114545,
			}
			if err := ml.CompareTestTensorSkippable(!onlyFirstLayer, expectedNextTokenOnlyFirstLayer, expectedNextTokenOnlyFirstLayerSize, actualNextToken, common.THRESHOLD_EXACT, true); err != nil {
				t.Fatal(err)
			}
		}
		actualNextTokenId := TokenId(actualNextToken.Item().(int32))
		// Comment in original Python code: only replace token if prompt has already been generated
		existingToken, err := tokens.GetItem([]int{curPos})
		if err != nil {
			t.Fatal(err)
		}
		existingTokenId := TokenId(existingToken.(int32))
		if existingTokenId != llamaModel.Vocabulary.PadId {
			actualNextTokenId = existingTokenId
		}
		if err = tokens.SetItem([]int{curPos}, int32(actualNextTokenId)); err != nil {
			t.Fatal(err)
		}
		_, actualEosReached := llamaModel.Vocabulary.StopTokenIds[actualNextTokenId]

		if onlyFirstLayer && isFirstIteration {
			expectedEosReached := false
			if actualEosReached != expectedEosReached {
				t.Fatalf("expected eosReached %v, but got %v", expectedEosReached, actualEosReached)
			}
		}
		prevPos = curPos
		isFirstIteration = false
	}

	if onlyFirstLayer {
		expectedOutputTokenIds := []TokenId{114545, 80657, 20508, 21053, 71434}
		actualOutputTokenIds := make([]TokenId, 0)
		for i := minPromptLength; i < infContext.SequenceLength; i++ {
			tokenItem, err := tokens.GetItem([]int{i})
			if err != nil {
				t.Fatal(err)
			}
			actualOutputTokenIds = append(actualOutputTokenIds, TokenId(tokenItem.(int32)))
		}
		if !reflect.DeepEqual(expectedOutputTokenIds, actualOutputTokenIds) {
			t.Fatalf("expected actualOutputTokenIds %v, but got %v", expectedOutputTokenIds, actualOutputTokenIds)
		}
	}
}

func TestSimulatedOnlyFirstLayer(t *testing.T) {
	t.Setenv("test.timeout", "10m")
	testSimulatedInternal(t, true)
}

const runTestSimulatedFull = false

func TestSimulatedFull(t *testing.T) {
	if !runTestSimulatedFull {
		t.Skip("Skipping TestSimulatedFull because runTestSimulatedFull is set to false")
	}
	testSimulatedInternal(t, false)
}
