Repository navigation
Conversation
|
Warning Review limit reached
Next review available in: 41 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis change introduces a Zig tensor framework with C kernel bindings, computation graphs, automatic differentiation, neural-network layers, optimizers, graph fusion, JIT scheduling, build integration, and a training executable. ChangesTensor framework and training pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Tensor
participant Functions
participant Graph
participant Scheduler
participant Optimizer
Tensor->>Functions: Create tensor operations
Functions->>Graph: Build nodes and dependencies
Scheduler->>Graph: Execute forward and backward DAG
Graph->>Tensor: Produce outputs and gradients
Optimizer->>Tensor: Update parameters and clear gradients
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
🟠 Major comments (16)
src/root.zig-108-121 (1)
108-121: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftUnbounded arena growth in the training loop.
The single
arenais used for params, inputs, and all per-epoch temporaries, but it is only freed at program exit. Each iterationnet.call,loss_fn.call, andsched.backwardallocate fresh tensors, nodes, and gradient buffers into this arena with no reset, so memory grows linearly across all 20000 epochs and can exhaust memory on larger nets/longer runs. Split allocation: keep a long-lived allocator fornet/params/x/y, and use a separate per-epoch arena that isreset(.retain_capacity)at the end of each iteration for the forward/backward scratch tensors.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/root.zig` around lines 108 - 121, Update the training loop around the epoch iteration to use a separate per-epoch arena for scratch allocations from net.call, loss_fn.call, and sched.backward, while keeping the existing long-lived allocator for net, parameters, x, and y. Reset the per-epoch arena with .retain_capacity at the end of every iteration, ensuring optimizer and model state remain allocated by the long-lived allocator.src/scheduler/scheduler.zig-41-45 (1)
41-45: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftJIT cache is bypassed; fusion re-runs on every forward/backward.
optimize()callsfusion.optimizedirectly and never touchesself.jit, so the cachedJIT.compilegraph is never used even whenjit_modeis on and aJITis injected (as inroot.zig). Consequently fusion pattern-matching + allocation runs onself.graphevery epoch instead of once, defeating the JIT’s purpose and adding per-epoch allocations. Either route throughself.jit.?.compile(&self.graph)or drop the unusedjitfield.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/scheduler/scheduler.zig` around lines 41 - 45, The optimize method bypasses the injected JIT cache by calling fusion.optimize directly. When jit_mode is enabled, route graph optimization through self.jit.?.compile(&self.graph) so compilation and fusion occur through the cached JIT path; otherwise remove the unused jit field and its injection if caching is not intended.src/nn/activation.zig-17-35 (1)
17-35: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftSame overflow-prone tanh formula duplicated in two places. Both compute
tanh(x)as(e^{2x}-1)/(e^{2x}+1), which overflows toNaNforf32inputs once|x|is roughly>44(relevant for training withT = f32persrc/root.zig). Fix once with a numerically stable formulation and share it.
src/nn/activation.zig#L17-L35: replaceTanh.call's formula with a stable variant (or a native tanh kernel) and reuse from rnn.zig.src/nn/rnn.zig#L61-L67: replace the inline duplicate with a call to the fixed shared helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/activation.zig` around lines 17 - 35, The overflow-prone tanh calculation must be replaced with one numerically stable shared implementation. Update Tanh.call in src/nn/activation.zig:17-35 to use a stable tanh helper or native kernel, then update the inline duplicate in src/nn/rnn.zig:61-67 to call that shared implementation instead of recomputing the exponential formula.src/nn/loss.zig-75-89 (1)
75-89: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
term1/term2can becomeNaNvia0 * -infwhen a prediction saturates to exactly 0 or 1.
target * log_inputandneg_target * log_one_minusare each vulnerable to IEEE-754's0 * -inf = NaN: iftarget == 0andinputunderflows to0.0(solog_input == -inf), ortarget == 1andinputsaturates to1.0(solog_one_minus == -inf), the product isNaNeven though the mathematically-correct contribution is0. This is plausible withf32sigmoid outputs. Standard practice is to clamp log outputs (e.g. to>= -100) before combining.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/loss.zig` around lines 75 - 89, Update call in the binary cross-entropy loss to clamp log_input and log_one_minus to a finite lower bound before multiplying by target and neg_target. Ensure saturated inputs cannot produce 0 * -inf or NaN while preserving the existing loss reduction flow.src/nn/batchnorm.zig-48-61 (1)
48-61: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
training/momentum/running_mean/running_varare stored but never used.
call()always computes normalization from the current batch's mean/variance, ignoringself.training. There's no update ofrunning_mean/running_varwithself.momentum, and no path that uses the running statistics for inference (training=false) — nor any setter to toggle training/eval mode. As written, BatchNorm1d behaves identically during training and evaluation, which defeats its purpose and will produce different results than PyTorch/standard BatchNorm semantics at inference time.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/batchnorm.zig` around lines 48 - 61, Update BatchNorm’s call method to honor the training flag: during training, normalize with batch statistics and update running_mean and running_var using momentum; during inference, normalize with the stored running statistics without updating them. Add or use an appropriate setter/API to toggle training versus evaluation mode, while preserving gamma, beta, and epsilon handling.src/nn/flatten.zig-12-27 (1)
12-27: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
forward(x, alloc)breaks the common layer interface used elsewhere (call(x)).Every other layer in this cohort (
ReLU,Tanh,Sigmoid,BatchNorm1d,Dropout,Linear,RNN,Softmax) exposescall(self, x) !*Tensor(T).Flatteninstead exposesforward(self, x, alloc).Sequential.add's wrapper (src/nn/sequential.zigline 32) unconditionally calls.call(x)on the wrapped layer, soFlattencannot be composed intoSequentialas-is, and even outsideSequentialit's inconsistent with the framework's convention.Suggested fix
- pub fn forward(self: *`@This`(), x: *Tensor(T), alloc: std.mem.Allocator) !*Tensor(T) { + pub fn call(self: *`@This`(), x: *Tensor(T)) !*Tensor(T) { + const alloc = x.alloc;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/flatten.zig` around lines 12 - 27, Rename Flatten’s public forward method to call and align its signature with the common layer interface by obtaining the allocator from the layer or existing framework context rather than requiring an alloc argument. Preserve the current flattening behavior and ensure Sequential’s wrapper can invoke Flatten through call(x).src/nn/loss.zig-99-125 (1)
99-125: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftRework
BCEWithLogitsLossto compute the loss from logits directly
It currently appliessigmoid(input)and then callsBCELoss, so it keeps the samelog(x)/log(1 - x)instability and adds extra work. Use the stable logits form instead (e.g.max(x, 0) - x * target + log1p(exp(-abs(x)))).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/loss.zig` around lines 99 - 125, Update BCEWithLogitsLoss.call to compute the loss directly from input logits using a numerically stable formulation such as max(input, 0) - input * target + log1p(exp(-abs(input))), rather than constructing sigmoid tensors and delegating to BCELoss. Preserve the configured reduction behavior and return a Tensor(T) with the same shape and API contract.src/nn/activation.zig-19-29 (1)
19-29: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftPer-iteration activation tensors accumulate for the whole run. These
call()paths allocate temporary tensors (one_t,eps_ptr,mask,ones/two_t/margin_t) from the same long-lived arena used insrc/root.zig, and there’s no per-iteration cleanup path for them.Graph.deinit()only drops the node list, so these activations stay resident until the final arena teardown.
src/nn/activation.zigsrc/nn/batchnorm.zigsrc/nn/dropout.zigsrc/nn/loss.zigsrc/nn/rnn.zig🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/activation.zig` around lines 19 - 29, Prevent temporary activation tensors from accumulating across iterations by adding per-iteration cleanup for allocations in activation call paths: src/nn/activation.zig lines 19-29 and 39-48, src/nn/batchnorm.zig lines 54-56, src/nn/dropout.zig lines 14-27, src/nn/loss.zig lines 156-180, and src/nn/rnn.zig lines 63-64. Ensure tensors such as one_t, eps_ptr, mask, ones, two_t, and margin_t use an iteration-scoped allocation lifecycle, while preserving the existing computation results and leaving Graph.deinit() responsible only for graph nodes.src/nn/flatten.zig-23-26 (1)
23-26: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Flattenshould be a graph-tracked view, not a raw buffer alias.creatornever gets set here, so autograd stops at this layer, andout.data = x.datamakes the output share the input’s allocation. Route this through the existing op/node path instead of constructing the tensor manually.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/flatten.zig` around lines 23 - 26, Update Flatten to create its result through the existing operation/node graph-tracking path rather than manually allocating Tensor(T) and assigning out.data. Ensure the resulting view records x as its creator and preserves autograd propagation without introducing a separate buffer alias; reuse the existing flatten op/node mechanism visible in this module.build.zig-34-40 (1)
34-40: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winAvoid hardcoded relative paths for system dependencies.
Hardcoding the CUDA path as
../../../../opt/cuda/includemakes the build brittle and prone to failure on other machines where CUDA might be installed in standard locations like/usr/local/cuda/includeorC:\Program Files\NVIDIA GPU Computing Toolkit.Consider using an environment variable like
CUDA_PATHor standard search paths to locate the CUDA headers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@build.zig` around lines 34 - 40, Update the CUDA configuration in the build script around cuda_bindings to resolve the CUDA include directory from CUDA_PATH or supported standard installation locations instead of hardcoding ../../../../opt/cuda/include. Use the resolved directory for both root_source_file and addIncludePath, and report a clear build error when no valid CUDA installation is found.src/ops/functions.zig-154-158 (2)
154-158: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMissing
keepdimhandling when reducing all dimensions.If
dimis null (meaning reduce over all dimensions) andkeepdimistrue, the output shape should preserve the rank of the input tensor by returning a shape of[1, 1, ..., 1]with a length equal toshape.len. The current logic unconditionally returns a 1D scalar shape[1], which can cause downstream broadcasting errors.🐛 Proposed fix
} else { - const out_shape = try allocator.alloc(u64, 1); - out_shape[0] = 1; - return out_shape; + if (keepdim) { + const out_shape = try allocator.alloc(u64, shape.len); + for (out_shape) |*s| s.* = 1; + return out_shape; + } else { + const out_shape = try allocator.alloc(u64, 1); + out_shape[0] = 1; + return out_shape; + } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ops/functions.zig` around lines 154 - 158, Update the all-dimensions reduction branch in the surrounding reduction-shape logic to honor keepdim: when keepdim is true, allocate an output shape with length shape.len and fill every entry with 1; retain the existing one-element shape for keepdim false. Ensure the allocation and initialization use the existing allocator/error propagation conventions.
154-158: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMissing
keepdimhandling when reducing all dimensions.If
dimis null (reducing over all dimensions) andkeepdimistrue, the output shape should preserve the rank of the input tensor by returning a shape of[1, 1, ..., 1]with lengthshape.len. The current logic unconditionally returns a 1D scalar shape[1], which can cause downstream broadcasting errors.🐛 Proposed fix
} else { - const out_shape = try allocator.alloc(u64, 1); - out_shape[0] = 1; - return out_shape; + if (keepdim) { + const out_shape = try allocator.alloc(u64, shape.len); + for (out_shape) |*s| s.* = 1; + return out_shape; + } else { + const out_shape = try allocator.alloc(u64, 1); + out_shape[0] = 1; + return out_shape; + } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ops/functions.zig` around lines 154 - 158, Update the all-dim reduction branch in the relevant function in functions.zig to honor keepdim: when keepdim is true, allocate an output shape with length shape.len and fill every dimension with 1; retain the existing one-dimensional [1] shape when keepdim is false.src/graph.zig-31-41 (1)
31-41: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAllocation failure in topological sort is silently swallowed, dropping the node from the execution order.
self.nodes.append(self.alloc, node) catch {};discards the error. If the append fails,nodeis missing fromself.nodes, soforward()/backward()will skip it entirely and downstream nodes execute against stale/never-computed data with no error surfaced anywhere.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/graph.zig` around lines 31 - 41, Update topologicalSort to propagate allocation failures from self.nodes.append instead of swallowing them, changing its return type as needed and propagating the error through recursive calls and every caller. Ensure forward() and backward() surface the failure and do not continue with an incomplete execution order.src/graph.zig-49-53 (1)
49-53: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
forward()silently swallows per-node execution errors and keeps iterating.
_ = node.forward(self.alloc) catch {};ignores failures fromNode.forward(e.g. allocation failures forc_inputs). Execution continues to subsequent nodes even though the current node's output buffer may not have been written, silently propagating incorrect/uninitialized data through the rest of the graph.🛡️ Proposed fix
-pub fn forward(self: *Self) void { +pub fn forward(self: *Self) !void { for (self.nodes.items) |node| { - _ = node.forward(self.alloc) catch {}; + _ = try node.forward(self.alloc); } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/graph.zig` around lines 49 - 53, Update Graph’s forward method to stop silently swallowing Node.forward failures: propagate the error from each node instead of catching and ignoring it, and adjust the method’s return type accordingly. Ensure iteration stops immediately when a node execution fails so later nodes cannot consume invalid output.src/tensor.zig-56-62 (1)
56-62: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUse of
assertfor public API input validation is unsafe in release builds.
assert(data.len == n)compiles to a no-op inReleaseFast/ReleaseSmall. Ifdata.len != n,allocator.dupe(T, data)allocates a buffer sized todata.lenwhileshape/stridesclaimnelements — any kernel that iterates up tonumel(shape)will read/write out of bounds on the underlying buffer.🛡️ Proposed fix
pub fn fromData(allocator: std.mem.Allocator, shape: []const u64, data: []const T, requires_grad: bool) !Self { var t = try init(allocator, shape, requires_grad); const n = numel(&t.shape, t.ndim); - assert(data.len == n); + if (data.len != n) return error.ShapeMismatch; t.data = try allocator.dupe(T, data);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tensor.zig` around lines 56 - 62, Replace the `assert(data.len == n)` check in `fromData` with release-safe public API validation that returns an appropriate error when `data.len` differs from the shape’s `n` elements. Only duplicate the data after validation, preserving the existing initialization and return behavior for valid input.src/tensor.zig-89-101 (1)
89-101: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winErrors are silently swallowed, and a pre-existing
self.gradleaks on repeatedbackward()calls.
self.alloc.create(...) catch return;andTensor(T).ones(...) catch return;abort with zero indication of failure — the caller has no way to know gradients were never computed, unlike the last step which at least prints on error.- If
self.gradis already non-null (e.g.backward()called a second time on the same tensor, common across training-loop iterations), it's overwritten atself.grad = grad;without freeing the previous grad tensor, leaking memory every call.🛡️ Proposed fix
pub fn backward(self: *Self) void { const creator = self.creator orelse return; self.realize(); - const grad = self.alloc.create(Tensor(T)) catch return; - grad.* = Tensor(T).ones(self.alloc, self.shape[0..self.ndim], false) catch return; + if (self.grad) |old| { + old.deinit(); + self.alloc.destroy(old); + } + const grad = self.alloc.create(Tensor(T)) catch |err| { + std.debug.print("Error allocating grad: {any}\n", .{err}); + return; + }; + grad.* = Tensor(T).ones(self.alloc, self.shape[0..self.ndim], false) catch |err| { + self.alloc.destroy(grad); + std.debug.print("Error initializing grad: {any}\n", .{err}); + return; + }; self.grad = grad;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tensor.zig` around lines 89 - 101, Update backward to report failures from alloc.create and Tensor(T).ones instead of silently returning, using the same error-reporting approach as graph.backward. Before assigning the new gradient at self.grad, release any existing gradient tensor so repeated backward calls do not leak memory; preserve the existing successful gradient-computation flow.
🟡 Minor comments (7)
src/optimizers/adam.zig-100-116 (1)
100-116: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winSilent parameter skips on transient allocation/op failure.
stepis declared!void, yet every op here usescatch continue, so an OOM or op failure silently drops the update for that parameter (and leavesp.datastale) instead of surfacing the error. This makes training failures invisible and hard to diagnose. Prefer propagating withtry(or at minimum logging) rather than swallowing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/optimizers/adam.zig` around lines 100 - 116, Update the parameter-update logic in step to propagate failures from scalarLike, functions.div, sqrtTensor, functions.add, functions.mul, functions.sub, and graph execution instead of using catch continue. Preserve successful parameter updates while allowing allocation or operation errors to reach the existing !void caller error path rather than leaving p.data unchanged silently.src/scheduler/fusion.zig-112-118 (1)
112-118: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winFused-out nodes are leaked when
allocis not an arena.
orderedRemovedrops the matchedNodepointers fromgraph.nodesbut never frees them. Tests use an arena so it's masked, but the JIT path callsfusion.optimizewithself.alloc, so each fusion leaks the replaced nodes (and their input slices). Consider deinit/free of removed nodes here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/scheduler/fusion.zig` around lines 112 - 118, Update the fusion logic around the orderedRemove loop to deinitialize and free each removed Node, including its input slice, before inserting fused at min. Use the existing alloc and Node cleanup conventions, while preserving the current removal order and insertion behavior.src/optimizers/sgd.zig-62-92 (1)
62-92: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
nesterovoption is accepted but never applied.
initexposesnesterovand stores it, butstepimplements only classic momentum; the flag has no effect. Either implement the Nesterov update or drop the option to avoid misleading callers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/optimizers/sgd.zig` around lines 62 - 92, Update the SGD `step` method to honor the stored `nesterov` option in the momentum path, implementing the intended Nesterov parameter update while preserving classic momentum behavior when it is disabled. If Nesterov cannot be supported, remove the `nesterov` option from `init` and its stored state so callers cannot enable an ineffective setting.src/scheduler/jit.zig-78-81 (1)
78-81: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove debug graph dump from
compile.These
std.debug.printcalls dump the full node list on every cache miss — a debugging artifact that will spam stderr in production.🧹 Proposed cleanup
const node_list = std.ArrayList(*Node(T)).fromOwnedSlice(captured_nodes); - std.debug.print("Compiled graph has {d} nodes:\n", .{node_list.items.len}); - for (node_list.items, 0..) |node, i| { - std.debug.print(" node {d}: op_type={any}\n", .{ i, node.op.op_type }); - }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/scheduler/jit.zig` around lines 78 - 81, Remove the two std.debug.print calls in compile that report the compiled graph and each node, while leaving the surrounding node-list compilation logic unchanged.src/scheduler/fusion.zig-126-135 (1)
126-135: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
matmul_bias_relushould not be wired throughmatmul_relu.The bias-fusion branch still selects
getFunction(.{ .fused = .matmul_relu }, 1, 0)and setsop_typeto.matmul_relu, even though it passes 3 inputs. If the dispatcher uses the fused op type to pick a kernel, this will route the bias path through the plain matmul+relu implementation and drop the bias term. Switch this branch to a dedicatedmatmul_bias_reluop/function if that path is meant to be fused.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/scheduler/fusion.zig` around lines 126 - 135, Update the bias-fusion branch around the fused node construction to use the dedicated matmul_bias_relu fused operation for both getFunction and op.op_type, while preserving its three-input fused_inputs wiring and existing parameters/output.src/nn/linear.zig-6-15 (1)
6-15: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winNon-Linux fallback seed is constant
When
opts.seedis null,defaultSeed()returns0on every non-Linux platform (and on Linux ifclock_gettimefails), so weight initialization becomes deterministic instead of random there. Use a cross-platform entropy source instead of falling back to zero.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/linear.zig` around lines 6 - 15, Update defaultSeed() to use a cross-platform entropy source whenever the Linux clock_gettime path is unavailable or fails, rather than returning the constant 0. Preserve the existing Linux timestamp-based seed when successful and ensure the fallback produces varying seed values across initializations.src/nn/loss.zig-145-155 (1)
145-155: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
cos_simdivides bynorm1*norm2with no epsilon guard — zero vectors produceinf/NaN.Suggested guard
- const cos_sim = try functions.div(T, dot, try functions.mul(T, norm1, norm2)); + const denom_eps = try alloc.create(Tensor(T)); + denom_eps.* = try Tensor(T).fromData(alloc, &.{1}, &.{1e-8}, false); + const denom = try functions.add(T, try functions.mul(T, norm1, norm2), denom_eps); + const cos_sim = try functions.div(T, dot, denom);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/loss.zig` around lines 145 - 155, Update the cosine-similarity calculation in the `call` method to guard the `norm1 * norm2` denominator with the established epsilon or zero-safe handling before `functions.div`. Ensure inputs with zero norms produce finite, defined output rather than `inf` or `NaN`, while preserving the existing calculation for nonzero vectors.
🧹 Nitpick comments (6)
src/optimizers/adam.zig (1)
6-15: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffExtract duplicated
scalarLike/sqrtTensorhelpers into a shared module. These helpers are copy-pasted across four optimizer files; a sharedoptimizers/common.zigwould remove the duplication and prevent them from diverging. Separately,sqrtTensorcomputes square root asexp(0.5 * log(x)), which is slower and less numerically robust than a direct sqrt op — worth centralizing so a single improved implementation benefits all optimizers.
src/optimizers/adam.zig#L6-L15: movescalarLikeandsqrtTensorto the shared module and import them.src/optimizers/adamw.zig#L6-L15: replace the local copies with the shared imports.src/optimizers/rmsprop.zig#L6-L15: replace the local copies with the shared imports.src/optimizers/sgd.zig#L6-L10: replace the localscalarLikewith the shared import.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/optimizers/adam.zig` around lines 6 - 15, Extract scalarLike and sqrtTensor into src/optimizers/common.zig, implementing sqrtTensor with the direct tensor square-root operation. In src/optimizers/adam.zig lines 6-15, src/optimizers/adamw.zig lines 6-15, and src/optimizers/rmsprop.zig lines 6-15, remove the local helpers and import the shared symbols; in src/optimizers/sgd.zig lines 6-10, remove scalarLike and use the shared import.src/nn/linear.zig (1)
31-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winManual uniform-init duplicates
Tensor(T).kaimingUniformused elsewhere (src/nn/rnn.zigline 18).Consider reusing the shared initializer for consistency instead of re-deriving bounds and looping manually here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/linear.zig` around lines 31 - 38, Update the weight initialization in the linear layer to reuse Tensor(T).kaimingUniform, matching the initializer used by the RNN implementation. Remove the manual bound calculation and rng-based data loop, while preserving the existing tensor shape, allocator, and seed behavior as supported by the shared initializer.src/nn/batchnorm.zig (1)
5-12: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
sqrtHelperis defined identically in two files. Both implementsqrtviaexp(0.5*log(x))with the exact same body — extract to a shared location (e.g.src/ops/functions.zig) instead of duplicating.
src/nn/batchnorm.zig#L5-L12: remove this copy, import the shared helper.src/nn/loss.zig#L5-L12: remove this copy, import the shared helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/nn/batchnorm.zig` around lines 5 - 12, The duplicate sqrtHelper implementations in src/nn/batchnorm.zig lines 5-12 and src/nn/loss.zig lines 5-12 should be replaced with one shared helper in the existing shared functions location, such as src/ops/functions.zig. Export the shared helper, update both batchnorm.zig and loss.zig to import and use it, and remove their local sqrtHelper definitions without changing behavior.src/ops/functions.zig (1)
9-27: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider refactoring
createNodeto acceptOpinstead ofOpType.Currently, reduction operations (
sum,mean,max,min) duplicate the memory allocation and node initialization logic found increateNodebecausecreateNodehardcodes theOp.params(likedim,keepdim,fval). ModifyingcreateNodeto accept a fully constructedOp(or an optionalparamsstruct) would allow all operations to reuse this centralized node creation logic, reducing duplication.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ops/functions.zig` around lines 9 - 27, Refactor createNode to accept a fully constructed Op rather than OpType, removing its hardcoded Op construction and preserving the supplied operation parameters. Update callers, including reduction operations such as sum, mean, max, and min, to construct and pass their Op values so they reuse the centralized allocation and node initialization logic.src/ops/reduce.zig (1)
11-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSame 1:1 forwarding-wrapper pattern as
element_wise.zig.See consolidated comment (anchored in
src/ops/element_wise.zig).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ops/reduce.zig` around lines 11 - 48, Replace the redundant 1:1 forwarding functions in getReduceFn with the corresponding c_api callbacks directly, following the established pattern in element_wise.zig. Update each sum, mean, max, and min Function mapping and remove the now-unused forward/backward wrapper functions.src/ops/element_wise.zig (1)
38-113: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEvery op requires a hand-written 1:1 forwarding wrapper pair — significant repetition that could grow unwieldy as ops are added.
See consolidated comment below (also applies to
src/ops/reduce.zig).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ops/element_wise.zig` around lines 38 - 113, Replace the hand-written forwarding wrapper pairs such as addForward/addBackward through negForward/negBackward with a reusable generic or comptime-generated dispatch mechanism that derives the forward and backward calls from each operation. Preserve the existing C ABI signatures and operation-specific behavior, and apply the same consolidation to the analogous wrappers in reduce.zig.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/graph.zig`:
- Around line 63-75: Update Graph.eql to compare each node’s complete operation
identity, including the specific op value inside op.op_type and all fields in
op.params, rather than only the outer union tag. Preserve the existing
comparisons for node count, input count, output dimensions, device, and shape so
structurally different operations or parameter values return false.
In `@src/nn/rnn.zig`:
- Around line 43-68: Update RNN.call to process the sequence recurrently:
iterate over the input timesteps, reuse the previous hidden state when computing
w_hh, and return the resulting sequence or final state according to the module’s
API. Ensure h_t is initialized once per call rather than reset each timestep,
and implement num_layers so configured layers are actually applied; preserve the
existing bias and activation behavior.
- Around line 70-77: Update the RNN parameters method to use the unmanaged
ArrayList lifecycle: initialize with .empty, pass alloc to each append, and
return an owned slice via toOwnedSlice(alloc). Ensure the temporary list is
deinitialized with deinit(alloc), including on error paths, rather than
returning list.items.
In `@src/nn/softmax.zig`:
- Around line 14-20: Update Softmax.call to pass the instance’s configured
self.dim and self.keepdim values to both functions.max and functions.sum instead
of null and false, preserving the selected reduction dimension and
keep-dimension behavior for per-row softmax computation.
In `@src/ops/functions.zig`:
- Line 178: Update each reduction output construction in src/ops/functions.zig
at lines 178-178, 204-204, 230-230, and 256-256 to pass x.requires_grad instead
of false to Tensor(T).zeros, preserving gradient tracking for all reduction
operations.
In `@src/ops/fused.zig`:
- Around line 24-28: Update matmulReluBackward so it explicitly panics instead
of silently discarding inputs, output, and params. Preserve the current
signature while making the unsupported fused backward path fail loudly until
gradient propagation is implemented.
In `@src/tensor.zig`:
- Around line 26-38: Update Tensor.init to validate that shape.len does not
exceed MAX_NDIM before copying into shape_arr, returning an appropriate error
for oversized shapes. Only perform the loop and compute_strides after this
check, preserving existing initialization for valid dimensions.
---
Major comments:
In `@build.zig`:
- Around line 34-40: Update the CUDA configuration in the build script around
cuda_bindings to resolve the CUDA include directory from CUDA_PATH or supported
standard installation locations instead of hardcoding
../../../../opt/cuda/include. Use the resolved directory for both
root_source_file and addIncludePath, and report a clear build error when no
valid CUDA installation is found.
In `@src/graph.zig`:
- Around line 31-41: Update topologicalSort to propagate allocation failures
from self.nodes.append instead of swallowing them, changing its return type as
needed and propagating the error through recursive calls and every caller.
Ensure forward() and backward() surface the failure and do not continue with an
incomplete execution order.
- Around line 49-53: Update Graph’s forward method to stop silently swallowing
Node.forward failures: propagate the error from each node instead of catching
and ignoring it, and adjust the method’s return type accordingly. Ensure
iteration stops immediately when a node execution fails so later nodes cannot
consume invalid output.
In `@src/nn/activation.zig`:
- Around line 17-35: The overflow-prone tanh calculation must be replaced with
one numerically stable shared implementation. Update Tanh.call in
src/nn/activation.zig:17-35 to use a stable tanh helper or native kernel, then
update the inline duplicate in src/nn/rnn.zig:61-67 to call that shared
implementation instead of recomputing the exponential formula.
- Around line 19-29: Prevent temporary activation tensors from accumulating
across iterations by adding per-iteration cleanup for allocations in activation
call paths: src/nn/activation.zig lines 19-29 and 39-48, src/nn/batchnorm.zig
lines 54-56, src/nn/dropout.zig lines 14-27, src/nn/loss.zig lines 156-180, and
src/nn/rnn.zig lines 63-64. Ensure tensors such as one_t, eps_ptr, mask, ones,
two_t, and margin_t use an iteration-scoped allocation lifecycle, while
preserving the existing computation results and leaving Graph.deinit()
responsible only for graph nodes.
In `@src/nn/batchnorm.zig`:
- Around line 48-61: Update BatchNorm’s call method to honor the training flag:
during training, normalize with batch statistics and update running_mean and
running_var using momentum; during inference, normalize with the stored running
statistics without updating them. Add or use an appropriate setter/API to toggle
training versus evaluation mode, while preserving gamma, beta, and epsilon
handling.
In `@src/nn/flatten.zig`:
- Around line 12-27: Rename Flatten’s public forward method to call and align
its signature with the common layer interface by obtaining the allocator from
the layer or existing framework context rather than requiring an alloc argument.
Preserve the current flattening behavior and ensure Sequential’s wrapper can
invoke Flatten through call(x).
- Around line 23-26: Update Flatten to create its result through the existing
operation/node graph-tracking path rather than manually allocating Tensor(T) and
assigning out.data. Ensure the resulting view records x as its creator and
preserves autograd propagation without introducing a separate buffer alias;
reuse the existing flatten op/node mechanism visible in this module.
In `@src/nn/loss.zig`:
- Around line 75-89: Update call in the binary cross-entropy loss to clamp
log_input and log_one_minus to a finite lower bound before multiplying by target
and neg_target. Ensure saturated inputs cannot produce 0 * -inf or NaN while
preserving the existing loss reduction flow.
- Around line 99-125: Update BCEWithLogitsLoss.call to compute the loss directly
from input logits using a numerically stable formulation such as max(input, 0) -
input * target + log1p(exp(-abs(input))), rather than constructing sigmoid
tensors and delegating to BCELoss. Preserve the configured reduction behavior
and return a Tensor(T) with the same shape and API contract.
In `@src/ops/functions.zig`:
- Around line 154-158: Update the all-dimensions reduction branch in the
surrounding reduction-shape logic to honor keepdim: when keepdim is true,
allocate an output shape with length shape.len and fill every entry with 1;
retain the existing one-element shape for keepdim false. Ensure the allocation
and initialization use the existing allocator/error propagation conventions.
- Around line 154-158: Update the all-dim reduction branch in the relevant
function in functions.zig to honor keepdim: when keepdim is true, allocate an
output shape with length shape.len and fill every dimension with 1; retain the
existing one-dimensional [1] shape when keepdim is false.
In `@src/root.zig`:
- Around line 108-121: Update the training loop around the epoch iteration to
use a separate per-epoch arena for scratch allocations from net.call,
loss_fn.call, and sched.backward, while keeping the existing long-lived
allocator for net, parameters, x, and y. Reset the per-epoch arena with
.retain_capacity at the end of every iteration, ensuring optimizer and model
state remain allocated by the long-lived allocator.
In `@src/scheduler/scheduler.zig`:
- Around line 41-45: The optimize method bypasses the injected JIT cache by
calling fusion.optimize directly. When jit_mode is enabled, route graph
optimization through self.jit.?.compile(&self.graph) so compilation and fusion
occur through the cached JIT path; otherwise remove the unused jit field and its
injection if caching is not intended.
In `@src/tensor.zig`:
- Around line 56-62: Replace the `assert(data.len == n)` check in `fromData`
with release-safe public API validation that returns an appropriate error when
`data.len` differs from the shape’s `n` elements. Only duplicate the data after
validation, preserving the existing initialization and return behavior for valid
input.
- Around line 89-101: Update backward to report failures from alloc.create and
Tensor(T).ones instead of silently returning, using the same error-reporting
approach as graph.backward. Before assigning the new gradient at self.grad,
release any existing gradient tensor so repeated backward calls do not leak
memory; preserve the existing successful gradient-computation flow.
---
Minor comments:
In `@src/nn/linear.zig`:
- Around line 6-15: Update defaultSeed() to use a cross-platform entropy source
whenever the Linux clock_gettime path is unavailable or fails, rather than
returning the constant 0. Preserve the existing Linux timestamp-based seed when
successful and ensure the fallback produces varying seed values across
initializations.
In `@src/nn/loss.zig`:
- Around line 145-155: Update the cosine-similarity calculation in the `call`
method to guard the `norm1 * norm2` denominator with the established epsilon or
zero-safe handling before `functions.div`. Ensure inputs with zero norms produce
finite, defined output rather than `inf` or `NaN`, while preserving the existing
calculation for nonzero vectors.
In `@src/optimizers/adam.zig`:
- Around line 100-116: Update the parameter-update logic in step to propagate
failures from scalarLike, functions.div, sqrtTensor, functions.add,
functions.mul, functions.sub, and graph execution instead of using catch
continue. Preserve successful parameter updates while allowing allocation or
operation errors to reach the existing !void caller error path rather than
leaving p.data unchanged silently.
In `@src/optimizers/sgd.zig`:
- Around line 62-92: Update the SGD `step` method to honor the stored `nesterov`
option in the momentum path, implementing the intended Nesterov parameter update
while preserving classic momentum behavior when it is disabled. If Nesterov
cannot be supported, remove the `nesterov` option from `init` and its stored
state so callers cannot enable an ineffective setting.
In `@src/scheduler/fusion.zig`:
- Around line 112-118: Update the fusion logic around the orderedRemove loop to
deinitialize and free each removed Node, including its input slice, before
inserting fused at min. Use the existing alloc and Node cleanup conventions,
while preserving the current removal order and insertion behavior.
- Around line 126-135: Update the bias-fusion branch around the fused node
construction to use the dedicated matmul_bias_relu fused operation for both
getFunction and op.op_type, while preserving its three-input fused_inputs wiring
and existing parameters/output.
In `@src/scheduler/jit.zig`:
- Around line 78-81: Remove the two std.debug.print calls in compile that report
the compiled graph and each node, while leaving the surrounding node-list
compilation logic unchanged.
---
Nitpick comments:
In `@src/nn/batchnorm.zig`:
- Around line 5-12: The duplicate sqrtHelper implementations in
src/nn/batchnorm.zig lines 5-12 and src/nn/loss.zig lines 5-12 should be
replaced with one shared helper in the existing shared functions location, such
as src/ops/functions.zig. Export the shared helper, update both batchnorm.zig
and loss.zig to import and use it, and remove their local sqrtHelper definitions
without changing behavior.
In `@src/nn/linear.zig`:
- Around line 31-38: Update the weight initialization in the linear layer to
reuse Tensor(T).kaimingUniform, matching the initializer used by the RNN
implementation. Remove the manual bound calculation and rng-based data loop,
while preserving the existing tensor shape, allocator, and seed behavior as
supported by the shared initializer.
In `@src/ops/element_wise.zig`:
- Around line 38-113: Replace the hand-written forwarding wrapper pairs such as
addForward/addBackward through negForward/negBackward with a reusable generic or
comptime-generated dispatch mechanism that derives the forward and backward
calls from each operation. Preserve the existing C ABI signatures and
operation-specific behavior, and apply the same consolidation to the analogous
wrappers in reduce.zig.
In `@src/ops/functions.zig`:
- Around line 9-27: Refactor createNode to accept a fully constructed Op rather
than OpType, removing its hardcoded Op construction and preserving the supplied
operation parameters. Update callers, including reduction operations such as
sum, mean, max, and min, to construct and pass their Op values so they reuse the
centralized allocation and node initialization logic.
In `@src/ops/reduce.zig`:
- Around line 11-48: Replace the redundant 1:1 forwarding functions in
getReduceFn with the corresponding c_api callbacks directly, following the
established pattern in element_wise.zig. Update each sum, mean, max, and min
Function mapping and remove the now-unused forward/backward wrapper functions.
In `@src/optimizers/adam.zig`:
- Around line 6-15: Extract scalarLike and sqrtTensor into
src/optimizers/common.zig, implementing sqrtTensor with the direct tensor
square-root operation. In src/optimizers/adam.zig lines 6-15,
src/optimizers/adamw.zig lines 6-15, and src/optimizers/rmsprop.zig lines 6-15,
remove the local helpers and import the shared symbols; in
src/optimizers/sgd.zig lines 6-10, remove scalarLike and use the shared import.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c356dd62-3e01-4745-8eae-f8c11babbc19
⛔ Files ignored due to path filters (1)
plast/plast_core.cpython-314-x86_64-linux-gnu.sois excluded by!**/*.so
📒 Files selected for processing (36)
.gitignorebuild.zigsetup.pysrc/c_api.zigsrc/function.zigsrc/graph.zigsrc/nn/activation.zigsrc/nn/batchnorm.zigsrc/nn/dropout.zigsrc/nn/flatten.zigsrc/nn/linear.zigsrc/nn/loss.zigsrc/nn/mod.zigsrc/nn/rnn.zigsrc/nn/sequential.zigsrc/nn/softmax.zigsrc/node.zigsrc/op.zigsrc/ops/element_wise.zigsrc/ops/functions.zigsrc/ops/fused.zigsrc/ops/linear.zigsrc/ops/mod.zigsrc/ops/reduce.zigsrc/optimizer/mod.zigsrc/optimizers/adam.zigsrc/optimizers/adamw.zigsrc/optimizers/mod.zigsrc/optimizers/rmsprop.zigsrc/optimizers/sgd.zigsrc/root.zigsrc/scheduler/fusion.zigsrc/scheduler/jit.zigsrc/scheduler/scheduler.zigsrc/tensor.zigtest_cimport
| pub fn eql(self: *Self, other: *Self) bool { | ||
| if (self.nodes.items.len != other.nodes.items.len) return false; | ||
| for (self.nodes.items, other.nodes.items) |a, b| { | ||
| if (@as(u8, @intFromEnum(std.meta.activeTag(a.op.op_type))) != @as(u8, @intFromEnum(std.meta.activeTag(b.op.op_type)))) return false; | ||
| if (a.inputs.len != b.inputs.len) return false; | ||
| if (a.output.ndim != b.output.ndim) return false; | ||
| if (a.output.device != b.output.device) return false; | ||
| for (a.output.shape[0..a.output.ndim], b.output.shape[0..b.output.ndim]) |sa, sb| { | ||
| if (sa != sb) return false; | ||
| } | ||
| } | ||
| return true; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
eql() only compares the outer OpType union tag, not the specific operation or its params — different ops are reported equal.
std.meta.activeTag(a.op.op_type) returns which union field is active (element_wise/reduce/linear/fused), not the inner enum value. So a graph built with .element_wise = .add is considered equal to one built with .element_wise = .mul (or .reduce = .sum vs .reduce = .max) as long as shapes/device/input-count match. a.op.params/b.op.params (e.g. dim, keepdim, fval) are also never compared, so reduce.sum(dim=0) and reduce.sum(dim=1) are indistinguishable too. Given this is intended to detect structurally identical graphs (e.g. for the JIT cache mentioned in the PR's scheduler layer), this bug can cause a cached/compiled kernel for one op to be incorrectly reused for a completely different op, silently producing wrong results. Note the existing tests (lines 133-195) only compare graphs built with the identical op on both sides, so they don't exercise this gap.
🛡️ Proposed fix
- if (`@as`(u8, `@intFromEnum`(std.meta.activeTag(a.op.op_type))) != `@as`(u8, `@intFromEnum`(std.meta.activeTag(b.op.op_type)))) return false;
+ if (!std.meta.eql(a.op.op_type, b.op.op_type)) return false;
+ if (!std.meta.eql(a.op.params, b.op.params)) return false;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| pub fn eql(self: *Self, other: *Self) bool { | |
| if (self.nodes.items.len != other.nodes.items.len) return false; | |
| for (self.nodes.items, other.nodes.items) |a, b| { | |
| if (@as(u8, @intFromEnum(std.meta.activeTag(a.op.op_type))) != @as(u8, @intFromEnum(std.meta.activeTag(b.op.op_type)))) return false; | |
| if (a.inputs.len != b.inputs.len) return false; | |
| if (a.output.ndim != b.output.ndim) return false; | |
| if (a.output.device != b.output.device) return false; | |
| for (a.output.shape[0..a.output.ndim], b.output.shape[0..b.output.ndim]) |sa, sb| { | |
| if (sa != sb) return false; | |
| } | |
| } | |
| return true; | |
| } | |
| pub fn eql(self: *Self, other: *Self) bool { | |
| if (self.nodes.items.len != other.nodes.items.len) return false; | |
| for (self.nodes.items, other.nodes.items) |a, b| { | |
| if (!std.meta.eql(a.op.op_type, b.op.op_type)) return false; | |
| if (!std.meta.eql(a.op.params, b.op.params)) return false; | |
| if (a.inputs.len != b.inputs.len) return false; | |
| if (a.output.ndim != b.output.ndim) return false; | |
| if (a.output.device != b.output.device) return false; | |
| for (a.output.shape[0..a.output.ndim], b.output.shape[0..b.output.ndim]) |sa, sb| { | |
| if (sa != sb) return false; | |
| } | |
| } | |
| return true; | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/graph.zig` around lines 63 - 75, Update Graph.eql to compare each node’s
complete operation identity, including the specific op value inside op.op_type
and all fields in op.params, rather than only the outer union tag. Preserve the
existing comparisons for node count, input count, output dimensions, device, and
shape so structurally different operations or parameter values return false.
| pub fn call(self: *@This(), x: *Tensor(T)) !*Tensor(T) { | ||
| const alloc = self.alloc; | ||
| const batch_size = x.shape[0]; | ||
| const h_t = try alloc.create(Tensor(T)); | ||
| h_t.* = try Tensor(T).zeros(alloc, &.{ batch_size, self.hidden_size }, x.requires_grad); | ||
|
|
||
| const ih = try functions.matmul(T, x, self.w_ih); | ||
| const hh = try functions.matmul(T, h_t, self.w_hh); | ||
| var pre = try functions.add(T, ih, hh); | ||
| if (self.b_ih) |b_ih| { | ||
| const with_b_ih = try functions.add(T, pre, b_ih); | ||
| if (self.b_hh) |b_hh| { | ||
| pre = try functions.add(T, with_b_ih, b_hh); | ||
| } else { | ||
| pre = with_b_ih; | ||
| } | ||
| } | ||
|
|
||
| const two = try functions.add(T, pre, pre); | ||
| const e2x = try functions.exp(T, two); | ||
| const one_t = try alloc.create(Tensor(T)); | ||
| one_t.* = try Tensor(T).ones(alloc, pre.shape[0..pre.ndim], false); | ||
| const num = try functions.sub(T, e2x, one_t); | ||
| const den = try functions.add(T, e2x, one_t); | ||
| return functions.div(T, num, den); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
RNN.call doesn't actually recur — h_t is always zero, making w_hh structurally dead, and num_layers is unused.
h_t is initialized to zeros at every call (line 47) and there is no loop over a sequence dimension that would feed a previous hidden state back in. hh = matmul(h_t, self.w_hh) therefore always evaluates to the zero tensor regardless of input, meaning w_hh never contributes to the output — this reduces to a single Linear+tanh layer on x, not a recurrent cell. num_layers is stored but never used to stack layers. As written, this module cannot process sequences or maintain state across timesteps.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/nn/rnn.zig` around lines 43 - 68, Update RNN.call to process the sequence
recurrently: iterate over the input timesteps, reuse the previous hidden state
when computing w_hh, and return the resulting sequence or final state according
to the module’s API. Ensure h_t is initialized once per call rather than reset
each timestep, and implement num_layers so configured layers are actually
applied; preserve the existing bias and activation behavior.
| pub fn parameters(self: *@This(), alloc: std.mem.Allocator) ![]*Tensor(T) { | ||
| var list = std.ArrayList(*Tensor(T)).init(alloc); | ||
| try list.append(alloc, self.w_ih); | ||
| try list.append(alloc, self.w_hh); | ||
| if (self.b_ih) |b| try list.append(alloc, b); | ||
| if (self.b_hh) |b| try list.append(alloc, b); | ||
| return list.items; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the project's declared Zig toolchain version to verify ArrayList API expectations.
fd '.zigversion|build.zig.zon' | xargs -I{} cat {} 2>/dev/null
rg -n 'minimum_zig_version|zig_version' build.zig.zon build.zig 2>/dev/nullRepository: Vixel2006/plast
Length of output: 153
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== files ==\n'
git ls-files | rg '(^|/)(build\.zig(\.zon)?|\.zigversion|src/nn/(rnn|sequential)\.zig)$'
printf '\n== zig version config ==\n'
for f in build.zig build.zig.zon .zigversion; do
if [ -f "$f" ]; then
echo "--- $f ---"
cat -n "$f"
fi
done
printf '\n== outlines ==\n'
for f in src/nn/rnn.zig src/nn/sequential.zig; do
if [ -f "$f" ]; then
echo "--- outline $f ---"
ast-grep outline "$f" --view expanded || true
fi
done
printf '\n== relevant lines in src/nn/rnn.zig ==\n'
sed -n '1,140p' src/nn/rnn.zig
printf '\n== relevant lines in src/nn/sequential.zig ==\n'
sed -n '1,140p' src/nn/sequential.zigRepository: Vixel2006/plast
Length of output: 11382
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== repository files ==' >&2
git ls-files | rg '(^|/)(build\.zig(\.zon)?|\.zigversion|src/nn/(rnn|sequential)\.zig)$' || true
for f in build.zig build.zig.zon .zigversion src/nn/rnn.zig src/nn/sequential.zig; do
if [ -f "$f" ]; then
echo
echo "===== $f ====="
cat -n "$f"
fi
done
echo
echo '== search for ArrayList usage =='
rg -n 'ArrayList|toOwnedSlice|\.empty|\.init\(' src/nn build.zig build.zig.zon .zigversion || trueRepository: Vixel2006/plast
Length of output: 14685
🌐 Web query:
Zig 0.15 std.ArrayList unmanaged .empty init allocator append toOwnedSlice documentation
💡 Result:
In Zig 0.15, the std.ArrayList type is an alias for what was previously known as std.ArrayListUnmanaged [1][2]. The structure no longer stores an allocator internally, making it smaller and requiring the allocator to be passed explicitly to every mutating operation [3][2][4]. Initialization You can initialize an ArrayList using the empty struct literal .{} [5][4]. var list: std.ArrayList(i32) =.{}; Operations requiring an allocator Because the allocator is not stored within the list, you must provide it to any method that modifies the list's memory or manages its lifetime [5][6]. - Deinitialization: Call list.deinit(allocator) to release all allocated memory [5][6]. - Appending: Call list.append(allocator, value) or list.appendSlice(allocator, items) [5][6]. - Conversion: toOwnedSlice(allocator) returns a slice containing the list's items, transfers ownership to the caller, and resets the list [3][7][5]. Example usage const std = @import("std"); var list: std.ArrayList(i32) =.{}; defer list.deinit(allocator); try list.append(allocator, 42); try list.appendSlice(allocator, &[_]i32{ 1, 2, 3 }); const owned_slice = try list.toOwnedSlice(allocator); defer allocator.free(owned_slice); Note: The "managed" version of the list—which stores an allocator internally—is now available under std.array_list.Managed but is considered deprecated for future removal [8][9][2]. Always prefer the standard std.ArrayList (formerly Unmanaged) in 0.15+ code [2][4].
Citations:
- 1: https://github.com/ziglang/zig/blob/0.15.2/lib/std/std.zig
- 2: https://ziggit.dev/t/arraylist-and-allocator-updating-code-to-0-15/12167
- 3: https://github.com/ziglang/zig/blob/41502c6aa53a3da31b276c23c4db74db7d04796b/lib/std/array_list.zig
- 4: https://sngeth.com/zig/systems-programming/breaking-changes/2025/10/24/zig-0-15-migration-roadblocks/
- 5: https://gist.github.com/pmarreck/44d95e869036027f9edf332ce9a94583
- 6: https://github.com/zigcc/skills/blob/main/zig-0.15/SKILL.md
- 7: https://github.com/ziglang/zig/blob/0b47e69b7c0aedbc142400305cda86ef58b41656/lib/std/array_list.zig
- 8: https://ziglang.org/download/0.15.1/release-notes.html
- 9: std.ArrayList: make unmanaged the default ziglang/zig#24801
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n 'parameters\(' src/nn src | sed -n '1,200p'Repository: Vixel2006/plast
Length of output: 4203
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== src/nn/linear.zig =='
sed -n '60,140p' src/nn/linear.zig
echo
echo '== src/root.zig =='
sed -n '50,110p' src/root.zigRepository: Vixel2006/plast
Length of output: 4404
Use the unmanaged ArrayList API here (src/nn/rnn.zig:70-76) std.ArrayList(*Tensor(T)).init(alloc) is the wrong pattern for this codebase, and returning list.items hands back a non-owned buffer. Switch to .empty, deinit(alloc), and toOwnedSlice(alloc).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/nn/rnn.zig` around lines 70 - 77, Update the RNN parameters method to use
the unmanaged ArrayList lifecycle: initialize with .empty, pass alloc to each
append, and return an owned slice via toOwnedSlice(alloc). Ensure the temporary
list is deinitialized with deinit(alloc), including on error paths, rather than
returning list.items.
| pub fn call(_: *@This(), x: *Tensor(T)) !*Tensor(T) { | ||
| const max_val = try functions.max(T, x, null, false); | ||
| const shifted = try functions.sub(T, x, max_val); | ||
| const e = try functions.exp(T, shifted); | ||
| const sum = try functions.sum(T, e, null, false); | ||
| return functions.div(T, e, sum); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
self.dim/self.keepdim are configured but never passed to functions.max/functions.sum.
call hardcodes null and false, so Softmax always reduces over the entire tensor rather than the dimension chosen at init(). For any input with batch_size > 1, this computes one softmax distribution across the whole batch+class tensor instead of a per-row distribution over the class dimension — which is almost certainly not the intended/expected behavior for classification use.
Suggested fix
- pub fn call(_: *`@This`(), x: *Tensor(T)) !*Tensor(T) {
- const max_val = try functions.max(T, x, null, false);
+ pub fn call(self: *`@This`(), x: *Tensor(T)) !*Tensor(T) {
+ const max_val = try functions.max(T, x, self.dim, self.keepdim);
const shifted = try functions.sub(T, x, max_val);
const e = try functions.exp(T, shifted);
- const sum = try functions.sum(T, e, null, false);
+ const sum = try functions.sum(T, e, self.dim, self.keepdim);
return functions.div(T, e, sum);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| pub fn call(_: *@This(), x: *Tensor(T)) !*Tensor(T) { | |
| const max_val = try functions.max(T, x, null, false); | |
| const shifted = try functions.sub(T, x, max_val); | |
| const e = try functions.exp(T, shifted); | |
| const sum = try functions.sum(T, e, null, false); | |
| return functions.div(T, e, sum); | |
| } | |
| pub fn call(self: *`@This`(), x: *Tensor(T)) !*Tensor(T) { | |
| const max_val = try functions.max(T, x, self.dim, self.keepdim); | |
| const shifted = try functions.sub(T, x, max_val); | |
| const e = try functions.exp(T, shifted); | |
| const sum = try functions.sum(T, e, self.dim, self.keepdim); | |
| return functions.div(T, e, sum); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/nn/softmax.zig` around lines 14 - 20, Update Softmax.call to pass the
instance’s configured self.dim and self.keepdim values to both functions.max and
functions.sum instead of null and false, preserving the selected reduction
dimension and keep-dimension behavior for per-row softmax computation.
| }; | ||
|
|
||
| const out = try alloc.create(Tensor(T)); | ||
| out.* = try Tensor(T).zeros(alloc, out_shape, false); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Reduction operations hardcode requires_grad = false, breaking backpropagation.
When creating the output tensor for reduction operations, requires_grad is explicitly set to false instead of preserving the input's gradient requirement. This breaks the computational graph because gradients will not flow back through any reduction operation (such as calculating a loss).
src/ops/functions.zig#L178-L178: Passx.requires_gradinstead offalse.src/ops/functions.zig#L204-L204: Passx.requires_gradinstead offalse.src/ops/functions.zig#L230-L230: Passx.requires_gradinstead offalse.src/ops/functions.zig#L256-L256: Passx.requires_gradinstead offalse.
📍 Affects 1 file
src/ops/functions.zig#L178-L178(this comment)src/ops/functions.zig#L204-L204src/ops/functions.zig#L230-L230src/ops/functions.zig#L256-L256
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/ops/functions.zig` at line 178, Update each reduction output construction
in src/ops/functions.zig at lines 178-178, 204-204, 230-230, and 256-256 to pass
x.requires_grad instead of false to Tensor(T).zeros, preserving gradient
tracking for all reduction operations.
| fn matmulReluBackward(inputs: [*c]?*const C_Tensor, output: *const C_Tensor, params: KernelParams) void { | ||
| _ = inputs; | ||
| _ = output; | ||
| _ = params; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Silently dropping gradients in fused backward pass.
The backward pass for matmul_relu is currently a silent no-op. If a user triggers a backward pass on a graph containing this fused operation, it will quietly drop the gradients instead of propagating them, leading to incorrect parameter updates.
If fused operations are intended only for inference or the backward logic is not yet implemented, this function should panic rather than returning silently.
🐛 Proposed fix to prevent silent failures
fn matmulReluBackward(inputs: [*c]?*const C_Tensor, output: *const C_Tensor, params: KernelParams) void {
_ = inputs;
_ = output;
_ = params;
+ `@panic`("Backward pass for fused operations is not implemented");
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn matmulReluBackward(inputs: [*c]?*const C_Tensor, output: *const C_Tensor, params: KernelParams) void { | |
| _ = inputs; | |
| _ = output; | |
| _ = params; | |
| } | |
| fn matmulReluBackward(inputs: [*c]?*const C_Tensor, output: *const C_Tensor, params: KernelParams) void { | |
| _ = inputs; | |
| _ = output; | |
| _ = params; | |
| `@panic`("Backward pass for fused operations is not implemented"); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/ops/fused.zig` around lines 24 - 28, Update matmulReluBackward so it
explicitly panics instead of silently discarding inputs, output, and params.
Preserve the current signature while making the unsupported fused backward path
fail loudly until gradient propagation is implemented.
| pub fn init(allocator: std.mem.Allocator, shape: []const u64, requires_grad: bool) !Self { | ||
| const ndim: u64 = @intCast(shape.len); | ||
| var shape_arr: [MAX_NDIM]u64 = undefined; | ||
| for (shape, 0..) |s, i| shape_arr[i] = s; | ||
| const strides = compute_strides(&shape_arr, ndim); | ||
| return .{ | ||
| .alloc = allocator, | ||
| .ndim = ndim, | ||
| .shape = shape_arr, | ||
| .strides = strides, | ||
| .requires_grad = requires_grad, | ||
| }; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Missing bounds check allows out-of-bounds write into fixed-size shape_arr.
shape_arr is [MAX_NDIM]u64 (size 8), but the loop for (shape, 0..) |s, i| shape_arr[i] = s; never validates shape.len <= MAX_NDIM. Any caller passing more than 8 dimensions triggers a safety panic in Debug/ReleaseSafe and undefined behavior (stack corruption) in ReleaseFast, since init already returns !Self.
🛡️ Proposed fix
pub fn init(allocator: std.mem.Allocator, shape: []const u64, requires_grad: bool) !Self {
+ if (shape.len > MAX_NDIM) return error.TooManyDimensions;
const ndim: u64 = `@intCast`(shape.len);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| pub fn init(allocator: std.mem.Allocator, shape: []const u64, requires_grad: bool) !Self { | |
| const ndim: u64 = @intCast(shape.len); | |
| var shape_arr: [MAX_NDIM]u64 = undefined; | |
| for (shape, 0..) |s, i| shape_arr[i] = s; | |
| const strides = compute_strides(&shape_arr, ndim); | |
| return .{ | |
| .alloc = allocator, | |
| .ndim = ndim, | |
| .shape = shape_arr, | |
| .strides = strides, | |
| .requires_grad = requires_grad, | |
| }; | |
| } | |
| pub fn init(allocator: std.mem.Allocator, shape: []const u64, requires_grad: bool) !Self { | |
| if (shape.len > MAX_NDIM) return error.TooManyDimensions; | |
| const ndim: u64 = `@intCast`(shape.len); | |
| var shape_arr: [MAX_NDIM]u64 = undefined; | |
| for (shape, 0..) |s, i| shape_arr[i] = s; | |
| const strides = compute_strides(&shape_arr, ndim); | |
| return .{ | |
| .alloc = allocator, | |
| .ndim = ndim, | |
| .shape = shape_arr, | |
| .strides = strides, | |
| .requires_grad = requires_grad, | |
| }; | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/tensor.zig` around lines 26 - 38, Update Tensor.init to validate that
shape.len does not exceed MAX_NDIM before copying into shape_arr, returning an
appropriate error for oversized shapes. Only perform the loop and
compute_strides after this check, preserving existing initialization for valid
dimensions.
Description
In this PR, I have done the first step of our library port from c to zig. this move seamed nessecary for the vision of the project. this project should be production-ready for robotics use, python is not good for low-level tasks, c api is not a pleasent experience to do. so porting the project to zig will lead to a cleaner code, less boilerplate, the only thing that will be left from the previous c code will be the
kernels/code. which will be also removed in the future with the introduction of thecodegenmodule that I'm planning to implement for auto-gen kernel optimization.Summary by CodeRabbit
New Features
Improvements