cwenzi/neuroflow-cpp

1#include "test_framework.hpp"
2#include "neuroflow/swiglu.hpp"
3#include <cmath>
4
5using namespace neuroflow;
6
7TEST(SwiGLU, ConstructionDefaultFF) {
8    SwiGLUFFN ffn(256);
9    EXPECT_EQ(ffn.d_model_, 256u);
10    EXPECT_GT(ffn.d_ff_, 0u);
11    EXPECT_NE(ffn.d_ff_, 256u);
12}
13
14TEST(SwiGLU, IntermediateSizeComputation) {
15    SwiGLUFFN ffn_256(256);
16    EXPECT_EQ(ffn_256.d_ff_, 256u * 4);
17
18    SwiGLUFFN ffn_512(512);
19    EXPECT_EQ(ffn_512.d_ff_, 512u * 4);
20
21    SwiGLUFFN ffn_custom(64, 128);
22    EXPECT_EQ(ffn_custom.d_ff_, 128u);
23}
24
25TEST(SwiGLU, ForwardOutputShape) {
26    SwiGLUFFN ffn(64, 128);
27    Tensor x({4, 64}, QuantType::FP32);
28    float* xp = x.as_fp32();
29    for (size_t i = 0; i < x.numel(); ++i) xp[i] = 0.1f;
30
31    Tensor out = ffn.forward(x);
32    EXPECT_EQ(out.shape_.size(), 2u);
33    EXPECT_EQ(out.shape_[0], 4u);
34    EXPECT_EQ(out.shape_[1], 64u);
35}
36
37TEST(SwiGLU, ForwardNoNaN) {
38    SwiGLUFFN ffn(64, 128);
39    Tensor x({2, 64}, QuantType::FP32);
40    float* xp = x.as_fp32();
41    for (size_t i = 0; i < x.numel(); ++i) xp[i] = 0.5f;
42
43    Tensor out = ffn.forward(x);
44    const float* op = out.as_fp32();
45    for (size_t i = 0; i < out.numel(); ++i) {
46        EXPECT_FALSE(std::isnan(op[i]));
47        EXPECT_FALSE(std::isinf(op[i]));
48    }
49}
50
51TEST(SwiGLU, BackwardGradientsExist) {
52    SwiGLUFFN ffn(64, 128);
53    ffn.training_mode_ = true;
54    Tensor x({2, 64}, QuantType::FP32);
55    float* xp = x.as_fp32();
56    for (size_t i = 0; i < x.numel(); ++i) xp[i] = 0.5f;
57
58    Tensor out = ffn.forward(x);
59
60    Tensor grad({2, 64}, QuantType::FP32);
61    float* gp = grad.as_fp32();
62    for (size_t i = 0; i < grad.numel(); ++i) gp[i] = 1.0f;
63
64    auto grads = ffn.backward(grad);
65    EXPECT_GT(grads.w_gate_weight_grad.numel(), 0u);
66    EXPECT_GT(grads.w_down_weight_grad.numel(), 0u);
67    EXPECT_EQ(grads.input_grad.shape_[0], 2u);
68    EXPECT_EQ(grads.input_grad.shape_[1], 64u);
69}
70
71int main() { RUN_ALL_TESTS(); }
72