issue/428: accommodate the changes to c/gguf tests

src/infiniop-test/src/ops/rope.cpp

+#include "infiniop/ops/rope.h"
 #include "ops.hpp"
 #include "utils.hpp"
 #include <infinirt.h>
@@ -6,6 +7,8 @@
 
 namespace infiniop_test::rope {
 struct Test::Attributes {
+    infiniopRoPEAlgo_t algo;
+
     std::shared_ptr<Tensor> y;
     std::shared_ptr<Tensor> x;
     std::shared_ptr<Tensor> pos_ids;
@@ -21,7 +24,7 @@ std::shared_ptr<Test> Test::build(
     auto test = std::shared_ptr<Test>(new Test(rtol, atol));
     test->_attributes = new Attributes();
 
-    if (tensors.find("y") == tensors.end()
+    if (!check_names(attributes, Test::attribute_names()) || tensors.find("y") == tensors.end()
         || tensors.find("x") == tensors.end()
         || tensors.find("pos_ids") == tensors.end()
         || tensors.find("sin_table") == tensors.end()
@@ -30,6 +33,8 @@ std::shared_ptr<Test> Test::build(
         throw std::runtime_error("Invalid Test");
     }
 
+    test->_attributes->algo = *reinterpret_cast<infiniopRoPEAlgo_t *>(attributes["algo"].data());
+
     test->_attributes->y = tensors["y"];
     test->_attributes->x = tensors["x"];
     test->_attributes->pos_ids = tensors["pos_ids"];
@@ -43,6 +48,7 @@ std::shared_ptr<Test> Test::build(
 std::shared_ptr<infiniop_test::Result> Test::run(
     infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) {
     infiniopRoPEDescriptor_t op_desc;
+    infiniopRoPEAlgo_t algo = _attributes->algo;
     auto y = _attributes->y->to(device, device_id);
     auto x = _attributes->x->to(device, device_id);
     auto pos_ids = _attributes->pos_ids->to(device, device_id);
@@ -54,7 +60,8 @@ std::shared_ptr<infiniop_test::Result> Test::run(
                                           x->desc(),
                                           pos_ids->desc(),
                                           sin_table->desc(),
-                                          cos_table->desc()),
+                                          cos_table->desc(),
+                                          algo),
              return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor."));
 
     size_t workspace_size;
@@ -101,7 +108,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
 }
 
 std::vector<std::string> Test::attribute_names() {
-    return {};
+    return {"algo"};
 }
 
 std::vector<std::string> Test::tensor_names() {
@@ -120,6 +127,7 @@ std::string Test::toString() const {
     oss << "- pos_ids: " << _attributes->pos_ids->info() << std::endl;
     oss << "- sin_table: " << _attributes->sin_table->info() << std::endl;
     oss << "- cos_table: " << _attributes->cos_table->info() << std::endl;
+    oss << "- algo: " << _attributes->algo << std::endl;
     oss << std::scientific << std::setprecision(2);
     oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
     return oss.str();

test/infiniop-test/test_generate/testcases/rope.py

@@ -2,27 +2,48 @@ from ast import List
 import numpy as np
 import gguf
 from typing import List
-
+from enum import Enum
 
 from .. import InfiniopTestWriter, InfiniopTestCase, np_dtype_to_ggml, gguf_strides, contiguous_gguf_strides
 
+class Algorithm(Enum):
+    GPT_J = 0
+    GPT_NEOX = 1
 
-def rotary_embedding(t, sin, cos):
-    dh = t.shape[2] 
-    assert dh % 2 == 0, "Embedding dimension must be even."
-
-    t_even = t[..., 0::2]  # [seq_len, n_head, dh // 2]
-    t_odd = t[..., 1::2]  # [seq_len, n_head, dh // 2]
 
+def rotary_embedding(t, sin, cos, algo):
+    def _rope(sin, cos, t1, t2):
         cos = np.expand_dims(cos, axis=1)  # [seq_len, 1, dh // 2]
         sin = np.expand_dims(sin, axis=1)  # [seq_len, 1, dh // 2]
 
-    t_out_even = t_even * cos - t_odd * sin
-    t_out_odd = t_even * sin + t_odd * cos
+        t_out_1 = t1 * cos - t2 * sin
+        t_out_2 = t1 * sin + t2 * cos
+
+        return t_out_1, t_out_2
+
+
+    dh = t.shape[-1]
+    assert dh % 2 == 0, "Embedding dimension must be even."
 
     t_out = np.empty_like(t)
+
+    if algo == Algorithm.GPT_J.value:
+        t_even = t[..., 0::2]  # [seq_len, n_head, dh // 2]
+        t_odd = t[..., 1::2]  # [seq_len, n_head, dh // 2]
+
+        t_out_even, t_out_odd = _rope(sin, cos, t_even, t_odd)
+
         t_out[..., 0::2] = t_out_even
         t_out[..., 1::2] = t_out_odd
+    else:
+        half_dim = dh // 2   
+        t_first = t[..., :half_dim]
+        t_second = t[..., half_dim:]
+
+        t_out_first, t_out_second = _rope(sin, cos, t_first, t_second)
+
+        t_out[..., :half_dim] = t_out_first
+        t_out[..., half_dim:] = t_out_second
 
     return t_out
 
@@ -52,6 +73,7 @@ class RoPETestCase(InfiniopTestCase):
         pos_ids: np.ndarray,
         sin_table: np.ndarray,
         cos_table: np.ndarray,
+        algo: int,
     ):
         super().__init__("rope")
         self.y = y
@@ -63,10 +85,12 @@ class RoPETestCase(InfiniopTestCase):
         self.pos_ids = pos_ids
         self.sin_table = sin_table
         self.cos_table = cos_table
+        self.algo = algo
 
     def write_test(self, test_writer: "InfiniopTestWriter"):
         super().write_test(test_writer)
 
+        test_writer.add_int32(test_writer.gguf_key("algo"), self.algo)
         test_writer.add_tensor(
             test_writer.gguf_key("y"), self.y, raw_dtype=np_dtype_to_ggml(self.y.dtype)
         )
@@ -97,6 +121,7 @@ class RoPETestCase(InfiniopTestCase):
             self.x.astype(np.float64),
             self.sin_table.astype(np.float64),
             self.cos_table.astype(np.float64),
+            self.algo,
         )
         test_writer.add_tensor(
             test_writer.gguf_key("ans"), ans, raw_dtype=gguf.GGMLQuantizationType.F64
@@ -121,10 +146,17 @@ if __name__ == "__main__":
         ((3, 32, 128), (8000, 200, 1), (7000, 128, 1)),
     ]
 
+
+    _ALGO = [
+        Algorithm.GPT_J,
+        Algorithm.GPT_NEOX,
+    ]
+
     _TENSOR_DTYPES_ = [np.float16, np.float32]
     test_writer = InfiniopTestWriter("rope.gguf")
     test_cases = []
 
+    for algo in _ALGO:
         for dtype in _TENSOR_DTYPES_:
             for shape, stride_x, stride_y in _TEST_CASES_:
                 x = np.random.rand(*shape).astype(dtype)
@@ -141,6 +173,7 @@ if __name__ == "__main__":
                     pos_ids=pos_ids,
                     sin_table=sin_table,
                     cos_table=cos_table,
+                    algo=algo.value,
                 )
                 test_cases.append(test_case)
     test_writer.add_tests(test_cases)