Repository navigation
Fork/190 implied bias review fixes - #7
Conversation
jeethu
commented
Oct 4, 2026
jeethu
commented
- ☑️ I understand it is strictly prohibited to use AI to write PR description
- AI usage disclosure:
Codex review: the comment above implied_bias_kernels_exact ran to 17 lines and repeated the rounding contract. The contract stays in ops.h, the Python docstrings and the tests; the code comments keep only the short reason. Same for the other multi-line implied-bias comments in ops.cpp, the CPU backend and the Metal host code and kernels. Comment-only change.
Codex review: gather_qmm accepts a 0-d biases factor since 477f91a, but its docs still said biases must have the batch dimensions of w. State the scalar form in the ops.h comment and the Python docstring, in the same words as quantized_matmul and dequantize.
Codex review: where the bias is materialized (fp32 x over fp16 scales, 3/5/6-bit, gather_qmm, CUDA) `scales * f` was a differentiable multiply and the factor got a nonzero gradient, while the implied-bias kernel path threw in QuantizedMatmul::vjp. dequantize already gave a nonzero factor gradient on both paths (its GPU vjp runs the fallback). The factor is a constant everywhere now: - materialize_implied_bias applies stop_gradient to the factor, so the materialized paths and the dequantize fallback give zero. - QuantizedMatmul::vjp returns zeros for the factor instead of throwing, and jvp gives a zero tangent for it. The scales gradient is unchanged. A throw on every path would need a custom_vjp, which export cannot serialize. Documented in ops.h and the quantized_matmul docstring. Test: test_implied_bias_factor_grad checks vjp and grad wrt the factor for quantized_matmul, gather_qmm and dequantize (and jvp for quantized_matmul) on CPU and GPU, on the kernel and materialized paths.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d3445fa8c1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| } else if (arg == 3 && implied) { | ||
| // The implied-bias factor is a constant. | ||
| vjps.push_back(zeros_like(primals[3], stream())); |
There was a problem hiding this comment.
When an implied-bias factor uses the direct kernel path, this branch makes its first-order VJP zero, but the scale VJP below still constructs bias_c directly from primals[3]. Consequently, differentiating the scale gradient with respect to the factor produces a nonzero mixed derivative, even though the newly documented contract says no gradient flows to the factor; materialized paths do not have this problem because they apply stop_gradient. Stop the factor before using it to form the scale VJP as well.
Useful? React with 👍 / 👎.
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
