From ff21f0ac8b0f37be27c7f9849e364cacc1c4a106 Mon Sep 17 00:00:00 2001 From: adrien Date: Mon, 25 May 2026 18:18:35 +0200 Subject: [PATCH] Added all operation and test for TensorAlloc, missing slice --- src/TensorAlloc.zig | 744 ++++++++++++++++++++++++++++++-------------- 1 file changed, 503 insertions(+), 241 deletions(-) diff --git a/src/TensorAlloc.zig b/src/TensorAlloc.zig index c996ac8..eb4c93f 100644 --- a/src/TensorAlloc.zig +++ b/src/TensorAlloc.zig @@ -65,6 +65,19 @@ pub fn Tensor( return new; } + /// Convert N-D coords (row-major) to flat index — fully comptime. + /// Usage: Tensor.idx(.{row, col}) + pub inline fn idx(comptime coords: [rank]usize) usize { + comptime { + var flat: usize = 0; + for (0..rank) |i| { + if (coords[i] >= shape[i]) @compileError("idx: Coordinate out of bounds"); + flat += coords[i] * strides_arr[i]; + } + return flat; + } + } + /// Element-wise add. Dimensions must match; scales resolve to finer. /// RHS must have the same shape as self, or total == 1 (broadcast). pub fn add(self: *const Self, alloc: Allocator, rhs: anytype) !Tensor( @@ -470,6 +483,197 @@ pub fn Tensor( return !(try self.eqAll(alloc, other)); } + pub fn contract( + self: *const Self, + alloc: Allocator, + rhs: anytype, + comptime axis_a: usize, + comptime axis_b: usize, + ) !blk: { + const RhsType = @TypeOf(rhs); + if (!sh.isTensor(RhsType)) + @compileError("rhs can only be a Tensor "); + if (axis_a >= rank) @compileError("contract: axis_a out of bounds"); + if (axis_b >= RhsType.rank) @compileError("contract: axis_b out of bounds"); + if (shape[axis_a] != RhsType.shape[axis_b]) @compileError("contract: shape mismatch at contraction axes"); + + const sa = sh.shapeRemoveAxis(shape, axis_a); + const sb = sh.shapeRemoveAxis(RhsType.shape, axis_b); + const rs_raw = sh.shapeCat(&sa, &sb); + const rs: []const comptime_int = if (rs_raw.len == 0) &.{1} else &rs_raw; + break :blk Tensor( + T, + dims.add(RhsType.dims).argsOpt(), + sh.finerScales(Self, RhsType).argsOpt(), + rs, + ); + } { + const RhsType = @TypeOf(rhs); + const k: usize = comptime shape[axis_a]; // contraction dimension + + const sa = comptime sh.shapeRemoveAxis(shape, axis_a); + const sb = comptime sh.shapeRemoveAxis(RhsType.shape, axis_b); + const rs_raw = comptime sh.shapeCat(&sa, &sb); + const rs: []const comptime_int = comptime if (rs_raw.len == 0) &.{1} else &rs_raw; + + const ResultType = Tensor( + T, + dims.add(RhsType.dims).argsOpt(), + sh.finerScales(Self, RhsType).argsOpt(), + rs, + ); + + const SelfNorm = Tensor(T, dims.argsOpt(), sh.finerScales(Self, RhsType).argsOpt(), shape); + const OtherNorm = Tensor(T, RhsType.dims.argsOpt(), sh.finerScales(Self, RhsType).argsOpt(), RhsType.shape); + + // Normalize both tensors to matching, finer scales safely with allocations + const l = try self.to(alloc, SelfNorm); + defer l.deinit(alloc); + const r = try rhs.to(alloc, OtherNorm); + defer r.deinit(alloc); + + const ResVec = @Vector(ResultType.total, T); + + // FAST PATH: Dot Product + if (comptime rank == 1 and RhsType.rank == 1 and axis_a == 0 and axis_b == 0) { + const result_vec: ResVec = if (comptime !sh.isInt(T)) blk: { + break :blk @splat(@reduce(.Add, l.data.* * r.data.*)); + } else blk: { + const mul_arr: [total]T = l.data.* *| r.data.*; + var acc: T = 0; + for (mul_arr) |val| acc +|= val; + break :blk @splat(acc); + }; + + const vec_ptr = try alloc.create(ResVec); + vec_ptr.* = result_vec; + return ResultType{ .data = vec_ptr }; + } + + // --- ZERO-COST COERCION TO ARRAYS FOR RUNTIME INDEXING --- + const a_arr: [total]T = l.data.*; + const b_arr: [RhsType.total]T = r.data.*; + + // FAST PATH: 2D Matrix Multiplication + if (comptime rank == 2 and RhsType.rank == 2 and axis_a == 1 and axis_b == 0) { + const rows = shape[0]; + const cols = RhsType.shape[1]; + const inner = shape[1]; + + var res_arr: [ResultType.total]T = undefined; + + for (0..rows) |i| { + for (0..cols) |j| { + var acc: T = 0; + for (0..inner) |id| { + const a_flat = i * _strides[0] + id * _strides[1]; + const b_flat = id * RhsType.strides_arr[0] + j * RhsType.strides_arr[1]; + + if (comptime sh.isInt(T)) acc +|= a_arr[a_flat] *| b_arr[b_flat] else acc += a_arr[a_flat] * b_arr[b_flat]; + } + res_arr[i * cols + j] = acc; + } + } + + const vec_ptr = try alloc.create(ResVec); + vec_ptr.* = res_arr; + return ResultType{ .data = vec_ptr }; + } + + // FALLBACK PATH + const rs_raw_strides = comptime sh.shapeStrides(&rs_raw); + var result_arr: [ResultType.total]T = undefined; + + for (0..ResultType.total) |res_flat| { + const res_coords = sh.decodeFlatCoords(res_flat, rs_raw.len, rs_raw_strides); + + var a_free: [sa.len]usize = undefined; + for (0..sa.len) |i| a_free[i] = res_coords[i]; + var b_free: [sb.len]usize = undefined; + for (0..sb.len) |i| b_free[i] = res_coords[sa.len + i]; + + var acc: T = 0; + for (0..k) |ki| { + const a_coords = sh.insertAxis(rank, axis_a, ki, &a_free); + const b_coords = sh.insertAxis(RhsType.rank, axis_b, ki, &b_free); + const a_flat = sh.encodeFlatCoords(&a_coords, rank, _strides); + const b_flat = sh.encodeFlatCoords(&b_coords, RhsType.rank, RhsType.strides_arr); + + if (comptime sh.isInt(T)) acc +|= a_arr[a_flat] *| b_arr[b_flat] else acc += a_arr[a_flat] * b_arr[b_flat]; + } + result_arr[res_flat] = acc; + } + + const vec_ptr = try alloc.create(ResVec); + vec_ptr.* = result_arr; + return ResultType{ .data = vec_ptr }; + } + + /// 3D Cross Product. Only defined for Rank-1 tensors of length 3. + /// Result dimensions are the sum of input dimensions. + pub fn cross(self: *const Self, alloc: Allocator, rhs: anytype) !Tensor( + T, + dims.add(@TypeOf(rhs).dims).argsOpt(), + sh.finerScales(Self, @TypeOf(rhs)).argsOpt(), + &.{3}, + ) { + const RhsType = @TypeOf(rhs); + + if (!sh.isTensor(RhsType)) + @compileError("rhs can only be a Tensor "); + if (comptime rank != 1 or shape[0] != 3 or RhsType.rank != 1 or RhsType.shape[0] != 3) + @compileError("cross product is only defined for 3D vectors (rank-1, length 3)"); + + // Bring both to the same scale (e.g., mm vs m) + const p = try self.resolveScalePair(alloc, rhs); + defer p.deinit(alloc); + const l = p.l.data; + const r = p.r.data; + + var res: [3]T = undefined; + if (comptime sh.isInt(T)) { + res[0] = (l[1] *| r[2]) -| (l[2] *| r[1]); + res[1] = (l[2] *| r[0]) -| (l[0] *| r[2]); + res[2] = (l[0] *| r[1]) -| (l[1] *| r[0]); + } else { + res[0] = (l[1] * r[2]) - (l[2] * r[1]); + res[1] = (l[2] * r[0]) - (l[0] * r[2]); + res[2] = (l[0] * r[1]) - (l[1] * r[0]); + } + + return try .load(alloc, res); + } + + /// Sum of squared elements. Cheaper than length(); use for ordering. + pub fn lengthSqr(self: *const Self) T { + return @reduce(.Add, self.data.* * self.data.*); + } + + /// Euclidean length (L2 norm). + pub fn length(self: *const Self) T { + const sq = self.lengthSqr(); + if (comptime @typeInfo(T) == .int) { + const UnsignedT = @Int(.unsigned, @typeInfo(T).int.bits); + return @as(T, @intCast(std.math.sqrt(@as(UnsignedT, @intCast(sq))))); + } + return @sqrt(sq); + } + + /// Product of all elements. Result has shape {1}; dimension exponent * total. + pub fn product(self: *const Self, alloc: Allocator) !Tensor( + T, + dims.scale(@as(comptime_int, total)).argsOpt(), + scales.argsOpt(), + &.{1}, + ) { + return Tensor( + T, + dims.scale(@as(comptime_int, total)).argsOpt(), + scales.argsOpt(), + &.{1}, + ).splat(alloc, @reduce(.Mul, self.data.*)); + } + pub fn formatNumber( self: *const Self, writer: *std.Io.Writer, @@ -886,247 +1090,305 @@ test "TensorAlloc | Scalar Imperial mass scales" { try std.testing.expectApproxEqAbs(2.5, total.data[0], 1e-6); } -// // ─── Vector / Tensor tests ──────────────────────────────────────────────── -// -// test "TensorAlloc | Vector initiate" { -// const Meter4 = Tensor(f32, .{ .L = 1 }, .{}, &.{4}); -// const m = Meter4.splat(1); -// try std.testing.expect(m.data[0] == 1); -// try std.testing.expect(m.data[3] == 1); -// } -// -// test "TensorAlloc | Vector format" { -// const MeterPerSecondSq = Tensor(f32, .{ .L = 1, .T = -2 }, .{ .T = .n }, &.{3}); -// const KgMeterPerSecond = Tensor(f32, .{ .M = 1, .L = 1, .T = -1 }, .{ .M = .k }, &.{3}); -// -// const accel = MeterPerSecondSq.splat(9.81); -// const momentum = KgMeterPerSecond{ .data = .{ 43, 0, 11 } }; -// -// var buf: [64]u8 = undefined; -// var res = try std.fmt.bufPrint(&buf, "{d}", .{accel}); -// try std.testing.expectEqualStrings("(9.81, 9.81, 9.81)m.ns⁻²", res); -// -// res = try std.fmt.bufPrint(&buf, "{d:.2}", .{momentum}); -// try std.testing.expectEqualStrings("(43.00, 0.00, 11.00)m.kg.s⁻¹", res); -// } -// -// test "TensorAlloc | Vector Vec3 Init and Basic Arithmetic" { -// const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); -// -// const v_zero = Meter3.zero; -// try std.testing.expectEqual(0, v_zero.data[0]); -// try std.testing.expectEqual(0, v_zero.data[2]); -// -// const v_one = Meter3.one; -// try std.testing.expectEqual(1, v_one.data[0]); -// -// const v_def = Meter3.splat(5); -// try std.testing.expectEqual(5, v_def.data[2]); -// -// const v1 = Meter3{ .data = .{ 10, 20, 30 } }; -// const v2 = Meter3{ .data = .{ 2, 4, 6 } }; -// -// const added = v1.add(v2); -// try std.testing.expectEqual(12, added.data[0]); -// try std.testing.expectEqual(24, added.data[1]); -// try std.testing.expectEqual(36, added.data[2]); -// -// const subbed = v1.sub(v2); -// try std.testing.expectEqual(8, subbed.data[0]); -// try std.testing.expectEqual(16, subbed.data[1]); -// try std.testing.expectEqual(24, subbed.data[2]); -// -// const neg = v1.negate(); -// try std.testing.expectEqual(-10, neg.data[0]); -// try std.testing.expectEqual(-20, neg.data[1]); -// try std.testing.expectEqual(-30, neg.data[2]); -// } -// -// test "TensorAlloc | Vector Kinematics (scalar mul/div broadcast)" { -// const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); -// const Second1 = Tensor(i32, .{ .T = 1 }, .{}, &.{1}); -// -// const pos = Meter3{ .data = .{ 100, 200, 300 } }; -// const time = Second1.splat(10); -// -// const vel = pos.div(time); -// try std.testing.expectEqual(10, vel.data[0]); -// try std.testing.expectEqual(20, vel.data[1]); -// try std.testing.expectEqual(30, vel.data[2]); -// try std.testing.expectEqual(1, @TypeOf(vel).dims.get(.L)); -// try std.testing.expectEqual(-1, @TypeOf(vel).dims.get(.T)); -// -// const new_pos = vel.mul(time); -// try std.testing.expectEqual(100, new_pos.data[0]); -// try std.testing.expectEqual(0, @TypeOf(new_pos).dims.get(.T)); -// } -// -// test "TensorAlloc | Vector Element-wise Math and Scaling" { -// const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); -// -// const v1 = Meter3{ .data = .{ 10, 20, 30 } }; -// const v2 = Meter3{ .data = .{ 2, 5, 10 } }; -// const dv = v1.div(v2); -// try std.testing.expectEqual(5, dv.data[0]); -// try std.testing.expectEqual(4, dv.data[1]); -// try std.testing.expectEqual(3, dv.data[2]); -// try std.testing.expectEqual(0, @TypeOf(dv).dims.get(.L)); -// } -// -// test "TensorAlloc | Vector Conversions" { -// const KiloMeter3 = Tensor(i32, .{ .L = 1 }, .{ .L = .k }, &.{3}); -// const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); -// -// const v_km = KiloMeter3{ .data = .{ 1, 2, 3 } }; -// const v_m = v_km.to(Meter3); -// try std.testing.expectEqual(1000, v_m.data[0]); -// try std.testing.expectEqual(2000, v_m.data[1]); -// try std.testing.expectEqual(3000, v_m.data[2]); -// try std.testing.expectEqual(UnitScale.none, @TypeOf(v_m).scales.get(.L)); -// } -// -// test "TensorAlloc | Vector Length" { -// const MeterInt3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); -// const MeterFloat3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); -// -// const v_int = MeterInt3{ .data = .{ 3, 4, 0 } }; -// try std.testing.expectEqual(25, v_int.lengthSqr()); -// try std.testing.expectEqual(5, v_int.length()); -// -// const v_float = MeterFloat3{ .data = .{ 3.0, 4.0, 0.0 } }; -// try std.testing.expectApproxEqAbs(@as(f32, 25.0), v_float.lengthSqr(), 1e-4); -// try std.testing.expectApproxEqAbs(@as(f32, 5.0), v_float.length(), 1e-4); -// } -// -// test "TensorAlloc | Vector Comparisons" { -// const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); -// const KiloMeter3 = Tensor(f32, .{ .L = 1 }, .{ .L = .k }, &.{3}); -// -// const v1 = Meter3{ .data = .{ 1000.0, 500.0, 0.0 } }; -// const v2 = KiloMeter3{ .data = .{ 1.0, 0.5, 0.0 } }; -// const v3 = KiloMeter3{ .data = .{ 1.0, 0.6, 0.0 } }; -// -// try std.testing.expect(v1.eqAll(v2)); -// try std.testing.expect(v1.neAll(v3)); -// -// const higher = v3.gt(v1); -// try std.testing.expectEqual(false, higher[0]); -// try std.testing.expectEqual(true, higher[1]); -// try std.testing.expectEqual(false, higher[2]); -// -// const equal = v3.eq(v1); -// try std.testing.expectEqual(true, equal[0]); -// try std.testing.expectEqual(false, equal[1]); -// try std.testing.expectEqual(true, equal[2]); -// -// const low_eq = v1.lte(v3); -// try std.testing.expect(low_eq[0] and low_eq[1] and low_eq[2]); -// } -// -// test "TensorAlloc | Vector vs Scalar broadcast comparison" { -// const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); -// const KiloMeter1 = Tensor(f32, .{ .L = 1 }, .{ .L = .k }, &.{1}); -// -// const positions = Meter3{ .data = .{ 500.0, 1200.0, 3000.0 } }; -// const threshold = KiloMeter1.splat(1); // 1 km = 1000 m -// -// const exceeded = positions.gt(threshold); -// try std.testing.expectEqual(false, exceeded[0]); -// try std.testing.expectEqual(true, exceeded[1]); -// try std.testing.expectEqual(true, exceeded[2]); -// -// const Meter1 = Tensor(f32, .{ .L = 1 }, .{}, &.{1}); -// const exact = positions.eq(Meter1.splat(500)); -// try std.testing.expect(exact[0] == true); -// try std.testing.expect(exact[1] == false); -// } -// -// test "TensorAlloc | Vector contract — dot product (rank-1 * rank-1)" { -// const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); -// const Newton3 = Tensor(f32, .{ .M = 1, .L = 1, .T = -2 }, .{}, &.{3}); -// -// const pos = Meter3{ .data = .{ 10.0, 0.0, 0.0 } }; -// const force = Newton3{ .data = .{ 5.0, 5.0, 0.0 } }; -// -// const work = force.contract(pos, 0, 0); -// try std.testing.expectEqual(50.0, work.data[0]); -// try std.testing.expectEqual(1, @TypeOf(work).dims.get(.M)); -// try std.testing.expectEqual(2, @TypeOf(work).dims.get(.L)); -// try std.testing.expectEqual(-2, @TypeOf(work).dims.get(.T)); -// } -// -// test "TensorAlloc | Vector contract — matrix multiply (rank-2 * rank-2)" { -// const A = Tensor(f32, .{}, .{}, &.{ 2, 3 }); -// const B = Tensor(f32, .{}, .{}, &.{ 3, 2 }); -// -// const a = A{ .data = .{ 1, 2, 3, 4, 5, 6 } }; -// const b = B{ .data = .{ 7, 8, 9, 10, 11, 12 } }; -// -// const c = a.contract(b, 1, 0); -// try std.testing.expectEqual(58, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 0, 0 })]); -// try std.testing.expectEqual(64, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 0, 1 })]); -// try std.testing.expectEqual(139, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 1, 0 })]); -// try std.testing.expectEqual(154, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 1, 1 })]); -// } -// -// test "TensorAlloc | Vector Abs, Pow, Sqrt and Product" { -// const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); -// -// const v1 = Meter3{ .data = .{ -2.0, 3.0, -4.0 } }; -// const v_abs = v1.abs(); -// try std.testing.expectEqual(2.0, v_abs.data[0]); -// try std.testing.expectEqual(4.0, v_abs.data[2]); -// -// const vol = v_abs.product(); -// try std.testing.expectEqual(24.0, vol.data[0]); -// try std.testing.expectEqual(3, @TypeOf(vol).dims.get(.L)); -// -// const area_vec = v_abs.pow(2); -// try std.testing.expectEqual(4.0, area_vec.data[0]); -// try std.testing.expectEqual(16.0, area_vec.data[2]); -// try std.testing.expectEqual(2, @TypeOf(area_vec).dims.get(.L)); -// -// const sqrted = area_vec.sqrt(); -// try std.testing.expectEqual(2, sqrted.data[0]); -// try std.testing.expectEqual(4, sqrted.data[2]); -// try std.testing.expectEqual(1, @TypeOf(sqrted).dims.get(.L)); -// } -// -// test "TensorAlloc | Vector eq broadcast on dimensionless" { -// const DimLess3 = Tensor(i32, .{}, .{}, &.{3}); -// const v = DimLess3{ .data = .{ 1, 2, 3 } }; -// -// const eq_res = v.eq(DimLess3.splat(2)); -// try std.testing.expectEqual(false, eq_res[0]); -// try std.testing.expectEqual(true, eq_res[1]); -// try std.testing.expectEqual(false, eq_res[2]); -// } -// -// test "TensorAlloc | Tensor idx helper and matrix access" { -// const Mat3x3 = Tensor(f32, .{}, .{}, &.{ 3, 3 }); -// var m: Mat3x3 = Mat3x3.zero; -// m.data[Mat3x3.idx(.{ 0, 0 })] = 1.0; -// m.data[Mat3x3.idx(.{ 1, 1 })] = 2.0; -// m.data[Mat3x3.idx(.{ 2, 2 })] = 3.0; -// -// try std.testing.expectEqual(1.0, m.data[0]); -// try std.testing.expectEqual(2.0, m.data[4]); -// try std.testing.expectEqual(3.0, m.data[8]); -// try std.testing.expectEqual(0.0, m.data[1]); -// } -// -// test "TensorAlloc | Tensor strides_arr correctness" { -// const T1 = Tensor(f32, .{}, .{}, &.{3}); -// const T2 = Tensor(f32, .{}, .{}, &.{ 3, 4 }); -// const T3 = Tensor(f32, .{}, .{}, &.{ 2, 3, 4 }); -// -// try std.testing.expectEqual(1, T1.strides_arr[0]); -// try std.testing.expectEqual(4, T2.strides_arr[0]); -// try std.testing.expectEqual(1, T2.strides_arr[1]); -// try std.testing.expectEqual(12, T3.strides_arr[0]); -// try std.testing.expectEqual(4, T3.strides_arr[1]); -// try std.testing.expectEqual(1, T3.strides_arr[2]); -// } +// ─── Vector / Tensor tests ──────────────────────────────────────────────── + +test "TensorAlloc | Vector initiate" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter4 = Tensor(f32, .{ .L = 1 }, .{}, &.{4}); + const m = try Meter4.splat(alloc, 1); + try std.testing.expect(m.data[0] == 1); + try std.testing.expect(m.data[3] == 1); +} + +test "TensorAlloc | Vector format" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const MeterPerSecondSq = Tensor(f32, .{ .L = 1, .T = -2 }, .{ .T = .n }, &.{3}); + const KgMeterPerSecond = Tensor(f32, .{ .M = 1, .L = 1, .T = -1 }, .{ .M = .k }, &.{3}); + + const accel = try MeterPerSecondSq.splat(alloc, 9.81); + const momentum = try KgMeterPerSecond.load(alloc, &.{ 43, 0, 11 }); + + var buf: [64]u8 = undefined; + var res = try std.fmt.bufPrint(&buf, "{d}", .{accel}); + try std.testing.expectEqualStrings("(9.81, 9.81, 9.81)m.ns⁻²", res); + + res = try std.fmt.bufPrint(&buf, "{d:.2}", .{momentum}); + try std.testing.expectEqualStrings("(43.00, 0.00, 11.00)m.kg.s⁻¹", res); +} + +test "TensorAlloc | Vector Vec3 Init and Basic Arithmetic" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); + + const v_zero = try Meter3.splat(alloc, 0); + try std.testing.expectEqual(0, v_zero.data[0]); + try std.testing.expectEqual(0, v_zero.data[2]); + + const v_one = try Meter3.splat(alloc, 1); + try std.testing.expectEqual(1, v_one.data[0]); + + const v_def = try Meter3.splat(alloc, 5); + try std.testing.expectEqual(5, v_def.data[2]); + + const v1 = try Meter3.load(alloc, &.{ 10, 20, 30 }); + const v2 = try Meter3.load(alloc, &.{ 2, 4, 6 }); + + const added = try v1.add(alloc, v2); + try std.testing.expectEqual(12, added.data[0]); + try std.testing.expectEqual(24, added.data[1]); + try std.testing.expectEqual(36, added.data[2]); + + const subbed = try v1.sub(alloc, v2); + try std.testing.expectEqual(8, subbed.data[0]); + try std.testing.expectEqual(16, subbed.data[1]); + try std.testing.expectEqual(24, subbed.data[2]); + + const neg = try v1.negate(alloc); + try std.testing.expectEqual(-10, neg.data[0]); + try std.testing.expectEqual(-20, neg.data[1]); + try std.testing.expectEqual(-30, neg.data[2]); +} + +test "TensorAlloc | Vector Kinematics (scalar mul/div broadcast)" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); + const Second1 = Tensor(i32, .{ .T = 1 }, .{}, &.{1}); + + const pos = try Meter3.load(alloc, &.{ 100, 200, 300 }); + const time = try Second1.splat(alloc, 10); + + const vel = try pos.div(alloc, time); + try std.testing.expectEqual(10, vel.data[0]); + try std.testing.expectEqual(20, vel.data[1]); + try std.testing.expectEqual(30, vel.data[2]); + try std.testing.expectEqual(1, @TypeOf(vel).dims.get(.L)); + try std.testing.expectEqual(-1, @TypeOf(vel).dims.get(.T)); + + const new_pos = try vel.mul(alloc, time); + try std.testing.expectEqual(100, new_pos.data[0]); + try std.testing.expectEqual(0, @TypeOf(new_pos).dims.get(.T)); +} + +test "TensorAlloc | Vector Element-wise Math and Scaling" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); + + const v1 = try Meter3.load(alloc, &.{ 10, 20, 30 }); + const v2 = try Meter3.load(alloc, &.{ 2, 5, 10 }); + + const dv = try v1.div(alloc, v2); + try std.testing.expectEqual(5, dv.data[0]); + try std.testing.expectEqual(4, dv.data[1]); + try std.testing.expectEqual(3, dv.data[2]); + try std.testing.expectEqual(0, @TypeOf(dv).dims.get(.L)); +} + +test "TensorAlloc | Vector Conversions" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const KiloMeter3 = Tensor(i32, .{ .L = 1 }, .{ .L = .k }, &.{3}); + const Meter3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); + + const v_km = try KiloMeter3.load(alloc, &.{ 1, 2, 3 }); + const v_m = try v_km.to(alloc, Meter3); + + try std.testing.expectEqual(1000, v_m.data[0]); + try std.testing.expectEqual(2000, v_m.data[1]); + try std.testing.expectEqual(3000, v_m.data[2]); + try std.testing.expectEqual(UnitScale.none, @TypeOf(v_m).scales.get(.L)); +} + +test "TensorAlloc | Vector Length" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const MeterInt3 = Tensor(i32, .{ .L = 1 }, .{}, &.{3}); + const MeterFloat3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); + + const v_int = try MeterInt3.load(alloc, &.{ 3, 4, 0 }); + try std.testing.expectEqual(25, v_int.lengthSqr()); + try std.testing.expectEqual(5, v_int.length()); + + const v_float = try MeterFloat3.load(alloc, &.{ 3.0, 4.0, 0.0 }); + try std.testing.expectApproxEqAbs(@as(f32, 25.0), v_float.lengthSqr(), 1e-4); + try std.testing.expectApproxEqAbs(@as(f32, 5.0), v_float.length(), 1e-4); +} + +test "TensorAlloc | Vector Comparisons" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); + const KiloMeter3 = Tensor(f32, .{ .L = 1 }, .{ .L = .k }, &.{3}); + + const v1 = try Meter3.load(alloc, &.{ 1000.0, 500.0, 0.0 }); + const v2 = try KiloMeter3.load(alloc, &.{ 1.0, 0.5, 0.0 }); + const v3 = try KiloMeter3.load(alloc, &.{ 1.0, 0.6, 0.0 }); + + try std.testing.expect(try v1.eqAll(alloc, v2)); + try std.testing.expect(try v1.neAll(alloc, v3)); + + const higher = try v3.gt(alloc, v1); + try std.testing.expectEqual(false, higher[0]); + try std.testing.expectEqual(true, higher[1]); + try std.testing.expectEqual(false, higher[2]); + + const equal = try v3.eq(alloc, v1); + try std.testing.expectEqual(true, equal[0]); + try std.testing.expectEqual(false, equal[1]); + try std.testing.expectEqual(true, equal[2]); + + const low_eq = try v1.lte(alloc, v3); + try std.testing.expect(low_eq[0] and low_eq[1] and low_eq[2]); +} + +test "TensorAlloc | Vector vs Scalar broadcast comparison" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); + const KiloMeter1 = Tensor(f32, .{ .L = 1 }, .{ .L = .k }, &.{1}); + + const positions = try Meter3.load(alloc, &.{ 500.0, 1200.0, 3000.0 }); + const threshold = try KiloMeter1.splat(alloc, 1); // 1 km = 1000 m + + const exceeded = try positions.gt(alloc, threshold); + try std.testing.expectEqual(false, exceeded[0]); + try std.testing.expectEqual(true, exceeded[1]); + try std.testing.expectEqual(true, exceeded[2]); + + const Meter1 = Tensor(f32, .{ .L = 1 }, .{}, &.{1}); + const exact = try positions.eq(alloc, try Meter1.splat(alloc, 500)); + try std.testing.expect(exact[0] == true); + try std.testing.expect(exact[1] == false); +} + +test "TensorAlloc | Vector contract — dot product (rank-1 * rank-1)" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); + const Newton3 = Tensor(f32, .{ .M = 1, .L = 1, .T = -2 }, .{}, &.{3}); + + const pos = try Meter3.load(alloc, &.{ 10.0, 0.0, 0.0 }); + const force = try Newton3.load(alloc, &.{ 5.0, 5.0, 0.0 }); + + const work = try force.contract(alloc, pos, 0, 0); + try std.testing.expectEqual(50.0, work.data[0]); + try std.testing.expectEqual(1, @TypeOf(work).dims.get(.M)); + try std.testing.expectEqual(2, @TypeOf(work).dims.get(.L)); + try std.testing.expectEqual(-2, @TypeOf(work).dims.get(.T)); +} + +test "TensorAlloc | Vector contract — matrix multiply (rank-2 * rank-2)" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const A = Tensor(f32, .{}, .{}, &.{ 2, 3 }); + const B = Tensor(f32, .{}, .{}, &.{ 3, 2 }); + + const a = try A.load(alloc, &.{ 1, 2, 3, 4, 5, 6 }); + const b = try B.load(alloc, &.{ 7, 8, 9, 10, 11, 12 }); + + const c = try a.contract(alloc, b, 1, 0); + try std.testing.expectEqual(58, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 0, 0 })]); + try std.testing.expectEqual(64, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 0, 1 })]); + try std.testing.expectEqual(139, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 1, 0 })]); + try std.testing.expectEqual(154, c.data[Tensor(f32, .{}, .{}, &.{ 2, 2 }).idx(.{ 1, 1 })]); +} + +test "TensorAlloc | Vector Abs, Pow, Sqrt and Product" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Meter3 = Tensor(f32, .{ .L = 1 }, .{}, &.{3}); + + const v1 = try Meter3.load(alloc, &.{ -2.0, 3.0, -4.0 }); + const v_abs = try v1.abs(alloc); + try std.testing.expectEqual(2.0, v_abs.data[0]); + try std.testing.expectEqual(4.0, v_abs.data[2]); + + const vol = try v_abs.product(alloc); + try std.testing.expectEqual(24.0, vol.data[0]); + try std.testing.expectEqual(3, @TypeOf(vol).dims.get(.L)); + + const area_vec = try v_abs.pow(alloc, 2); + try std.testing.expectEqual(4.0, area_vec.data[0]); + try std.testing.expectEqual(16.0, area_vec.data[2]); + try std.testing.expectEqual(2, @TypeOf(area_vec).dims.get(.L)); + + const sqrted = try area_vec.sqrt(alloc); + try std.testing.expectEqual(2, sqrted.data[0]); + try std.testing.expectEqual(4, sqrted.data[2]); + try std.testing.expectEqual(1, @TypeOf(sqrted).dims.get(.L)); +} + +test "TensorAlloc | Vector eq broadcast on dimensionless" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const DimLess3 = Tensor(i32, .{}, .{}, &.{3}); + const v = try DimLess3.load(alloc, &.{ 1, 2, 3 }); + + const eq_res = try v.eq(alloc, try DimLess3.splat(alloc, 2)); + try std.testing.expectEqual(false, eq_res[0]); + try std.testing.expectEqual(true, eq_res[1]); + try std.testing.expectEqual(false, eq_res[2]); +} + +test "TensorAlloc | Tensor idx helper and matrix access" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + const Mat3x3 = Tensor(f32, .{}, .{}, &.{ 3, 3 }); + var m = try Mat3x3.splat(alloc, 0); + m.data[Mat3x3.idx(.{ 0, 0 })] = 1; + m.data[Mat3x3.idx(.{ 1, 1 })] = 2; + m.data[Mat3x3.idx(.{ 2, 2 })] = 3; + + try std.testing.expectEqual(1.0, m.data[0]); + try std.testing.expectEqual(2.0, m.data[4]); + try std.testing.expectEqual(3.0, m.data[8]); + try std.testing.expectEqual(0.0, m.data[1]); +} + +test "TensorAlloc | Tensor strides_arr correctness" { + const T1 = Tensor(f32, .{}, .{}, &.{3}); + const T2 = Tensor(f32, .{}, .{}, &.{ 3, 4 }); + const T3 = Tensor(f32, .{}, .{}, &.{ 2, 3, 4 }); + + try std.testing.expectEqual(1, T1.strides_arr[0]); + try std.testing.expectEqual(4, T2.strides_arr[0]); + try std.testing.expectEqual(1, T2.strides_arr[1]); + try std.testing.expectEqual(12, T3.strides_arr[0]); + try std.testing.expectEqual(4, T3.strides_arr[1]); + try std.testing.expectEqual(1, T3.strides_arr[2]); +} // // test "TensorAlloc | Slice 1D basic" { // const Vec = Tensor(i32, .{}, .{}, &.{5});