nx_tensor_test.nx source
↩ module page · 152 lines · 6048 B
1// nx_tensor_test.nx -- foundation tensor primitive shape verification.
2
3import "nx_syscalls.nx"
4import "nx_tier.nx"
5import "nx_tensor.nx"
6
7func main() -> nx_int {
8 let err: *i64 = (sys_mmap(8)) as *i64
9
10 // ===== Bad dtype rejected =====================================
11 let shape0: *i64 = (sys_mmap(16)) as *i64
12 shape0[0] = 4
13 let bad1: *NxTensor = nx_t_alloc(0 - 1, shape0, 1, err)
14 if (bad1 as nx_int) != 0 { return 1 }
15 if err[0] != NX_T_ERR_BAD_DTYPE { return 2 }
16
17 let bad2: *NxTensor = nx_t_alloc(NX_DT_N_KINDS, shape0, 1, err)
18 if (bad2 as nx_int) != 0 { return 3 }
19 if err[0] != NX_T_ERR_BAD_DTYPE { return 4 }
20
21 // ===== Gated dtype rejected with DTYPE_GATED ==================
22 let bad3: *NxTensor = nx_t_alloc(NX_DT_F32, shape0, 1, err)
23 if (bad3 as nx_int) != 0 { return 5 }
24 if err[0] != NX_T_ERR_DTYPE_GATED { return 6 }
25
26 let bad4: *NxTensor = nx_t_alloc(NX_DT_Q4_K, shape0, 1, err)
27 if err[0] != NX_T_ERR_DTYPE_GATED { return 7 }
28
29 // ===== ndim bounds ============================================
30 let bad5: *NxTensor = nx_t_alloc(NX_DT_I64, shape0, 0 - 1, err)
31 if err[0] != NX_T_ERR_BAD_NDIM { return 8 }
32 let bad6: *NxTensor = nx_t_alloc(NX_DT_I64, shape0, 9, err)
33 if err[0] != NX_T_ERR_BAD_NDIM { return 9 }
34
35 // ===== Bad shape (zero or negative) ==========================
36 shape0[0] = 0
37 let bad7: *NxTensor = nx_t_alloc(NX_DT_I64, shape0, 1, err)
38 if err[0] != NX_T_ERR_BAD_SHAPE { return 10 }
39 shape0[0] = 0 - 1
40 let bad8: *NxTensor = nx_t_alloc(NX_DT_I64, shape0, 1, err)
41 if err[0] != NX_T_ERR_BAD_SHAPE { return 11 }
42
43 // ===== Happy path: 1-D tensor of length 5 ====================
44 shape0[0] = 5
45 let t1: *NxTensor = nx_t_alloc(NX_DT_I64, shape0, 1, err)
46 if (t1 as nx_int) == 0 { return 20 }
47 if err[0] != NX_T_OK { return 21 }
48 if t1.dtype != NX_DT_I64 { return 22 }
49 if t1.ndim != 1 { return 23 }
50 if t1.numel != 5 { return 24 }
51 if t1.elt_bytes != 8 { return 25 }
52 if t1.storage_n != 40 { return 26 }
53 if t1.shape[0] != 5 { return 27 }
54 if t1.stride[0] != 1 { return 28 }
55 if nx_t_is_contiguous(t1) != 1 { return 29 }
56
57 // ===== Get/set + zero-fill ====================================
58 nx_t_fill_zero(t1)
59 let idx: *i64 = (sys_mmap(64)) as *i64
60 idx[0] = 0
61 if nx_t_get_i64(t1, idx) != 0 { return 30 }
62 nx_t_set_i64(t1, idx, 42)
63 if nx_t_get_i64(t1, idx) != 42 { return 31 }
64 idx[0] = 4
65 nx_t_set_i64(t1, idx, 1000)
66 if nx_t_get_i64(t1, idx) != 1000 { return 32 }
67
68 // ===== 2-D tensor: 3x4 =======================================
69 let shape2: *i64 = (sys_mmap(16)) as *i64
70 shape2[0] = 3
71 shape2[1] = 4
72 let t2: *NxTensor = nx_t_alloc(NX_DT_I64, shape2, 2, err)
73 if t2.numel != 12 { return 40 }
74 // Row-major: stride[0] = 4 (one row), stride[1] = 1
75 if t2.stride[0] != 4 { return 41 }
76 if t2.stride[1] != 1 { return 42 }
77 if nx_t_is_contiguous(t2) != 1 { return 43 }
78
79 nx_t_fill_zero(t2)
80 idx[0] = 1; idx[1] = 2
81 nx_t_set_i64(t2, idx, 77)
82 if nx_t_get_i64(t2, idx) != 77 { return 44 }
83 // Make sure other cells are still zero
84 idx[0] = 0; idx[1] = 0
85 if nx_t_get_i64(t2, idx) != 0 { return 45 }
86 idx[0] = 2; idx[1] = 3
87 if nx_t_get_i64(t2, idx) != 0 { return 46 }
88
89 // ===== Reshape happy path: 3x4 -> 4x3 ========================
90 let new_shape: *i64 = (sys_mmap(16)) as *i64
91 new_shape[0] = 4; new_shape[1] = 3
92 if nx_t_reshape(t2, new_shape, 2) != 0 { return 50 }
93 if t2.shape[0] != 4 { return 51 }
94 if t2.shape[1] != 3 { return 52 }
95 // Row-major: stride[0] = 3 (one row), stride[1] = 1
96 if t2.stride[0] != 3 { return 53 }
97 if t2.stride[1] != 1 { return 54 }
98 // Element at [1,2] in 3x4 layout was at linear offset 1*4+2 = 6.
99 // In 4x3 layout, linear offset 6 is at [2,0].
100 idx[0] = 2; idx[1] = 0
101 if nx_t_get_i64(t2, idx) != 77 { return 55 }
102
103 // ===== Reshape rejects bad numel ============================
104 new_shape[0] = 5
105 if nx_t_reshape(t2, new_shape, 2) != 0 - 1 { return 60 }
106 // shape unchanged after failed reshape
107 if t2.shape[0] != 4 { return 61 }
108
109 // ===== Permute: 3-D 2x3x4 -> permute to 4x2x3 ===============
110 let shape3: *i64 = (sys_mmap(24)) as *i64
111 shape3[0] = 2; shape3[1] = 3; shape3[2] = 4
112 let t3: *NxTensor = nx_t_alloc(NX_DT_I64, shape3, 3, err)
113 nx_t_fill_zero(t3)
114 // Put a marker at original [1, 2, 3]
115 idx[0] = 1; idx[1] = 2; idx[2] = 3
116 nx_t_set_i64(t3, idx, 999)
117
118 // Permutation [2, 0, 1] means: new[a, b, c] = old[b, c, a]
119 let perm: *i64 = (sys_mmap(24)) as *i64
120 perm[0] = 2; perm[1] = 0; perm[2] = 1
121 let v3: *NxTensor = nx_t_permute(t3, perm)
122 if (v3 as nx_int) == 0 { return 70 }
123 if v3.ndim != 3 { return 71 }
124 if v3.shape[0] != 4 { return 72 } // was shape[2]
125 if v3.shape[1] != 2 { return 73 } // was shape[0]
126 if v3.shape[2] != 3 { return 74 } // was shape[1]
127 // Permuted view typically NOT contiguous
128 if nx_t_is_contiguous(v3) != 0 { return 75 }
129 // The 999 marker: original [1,2,3] reads as view [3,1,2]
130 idx[0] = 3; idx[1] = 1; idx[2] = 2
131 if nx_t_get_i64(v3, idx) != 999 { return 76 }
132
133 // Permute rejects out-of-range axis
134 perm[0] = 5
135 let bad_v: *NxTensor = nx_t_permute(t3, perm)
136 if (bad_v as nx_int) != 0 { return 80 }
137 // Permute rejects duplicate axis
138 perm[0] = 1; perm[1] = 1; perm[2] = 2
139 let dup_v: *NxTensor = nx_t_permute(t3, perm)
140 if (dup_v as nx_int) != 0 { return 81 }
141
142 // ===== Helper coverage =======================================
143 if nx_dt_is_valid(NX_DT_I64) != 1 { return 90 }
144 if nx_dt_is_valid(NX_DT_F32) != 1 { return 91 }
145 if nx_dt_is_valid(NX_DT_N_KINDS) != 0 { return 92 }
146 if nx_dt_is_implemented(NX_DT_I64) != 1 { return 93 }
147 if nx_dt_is_implemented(NX_DT_F32) != 0 { return 94 } // gated
148 if nx_t_verdict_is_valid(NX_T_OK) != 1 { return 95 }
149 if nx_t_verdict_is_valid(NX_T_ERR_N_VERDICTS) != 0 { return 96 }
150
151 return 0
152}