reshape support batch to/from dim#6784
Conversation
|
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6784 +/- ##
===========================================
- Coverage 95.20% 93.08% -2.12%
===========================================
Files 946 548 -398
Lines 414000 230323 -183677
===========================================
- Hits 394128 214385 -179743
+ Misses 19872 15938 -3934 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 86ad7713e4
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a71cf4d120
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9884bffd8a
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aca6df296b
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 83f81b82df
ℹ️ 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".
| shape = op->params.at("shape").ai; | ||
| } | ||
|
|
||
| const int batch_size = op->inputs[0]->shape[batch_index]; |
There was a problem hiding this comment.
Guard reshape batch-axis shape lookup
When a batch-indexed tensor reaches Tensor.reshape with an unknown or omitted shape, this new unconditional lookup indexes op->inputs[0]->shape[batch_index] even though input_rank0 may be 0. Other code in this pass explicitly handles empty shapes, so converting a dynamic-shape model with a reshape after a known batch op can crash during solve_batch_index instead of falling back; the backward reshape path now has the same unchecked output-shape lookup.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5797ec95f3
ℹ️ 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".
| int sa = axes[si]; | ||
| if (shape.empty()) | ||
| continue; | ||
| if (sa < 0 || sa > (int)shape.size()) | ||
| continue; | ||
|
|
||
| shape.insert(shape.begin() + sa, 1); |
There was a problem hiding this comment.
Preserve full-axis positions for slice_copy select reshapes
When a selected slice_copy axis is after a hidden batch axis, axes[si] has already been decremented for the ncnn SliceCopy parameters, but in->shape here is still the full torch shape including the batch dimension. For example, updating x[:, idx, :] with shapes [N,C,H] and [N,H] at batch_index == 0 changes axis 1 to 0, so this inserts the singleton as [1,N,H] instead of [N,1,H], causing the value tensor to be reshaped into the wrong layout before the copy.
Useful? React with 👍 / 👎.
| const int batch_index = op->inputs[0]->params["__batch_index"].i; | ||
| if (batch_index != 233) | ||
| fprintf(stderr, "sum along batch axis is not supported yet\n"); | ||
|
|
||
| op->params["0"] = 0; | ||
| op->params["1"] = 1; |
There was a problem hiding this comment.
Do not lower reduce-all on batched tensors
For torch.sum(x) with a hidden batch axis, this warning is followed by emitting reduce_all=1, but ncnn layers with support_batch == false are invoked once per Mat.n batch, so the batch dimension is never included in the reduction. The converted model returns one reduced value per batch instead of the single scalar that PyTorch produces; the same pattern affects the other no-dim reduction rewrites added here.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c642fe0241
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1d7270dfb4
ℹ️ 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".
| const int batch_index = op->inputs[0]->params["__batch_index"].i; | ||
| if (batch_index != 233) | ||
| fprintf(stderr, "mean along batch axis is not supported yet\n"); |
There was a problem hiding this comment.
Reject mean over the hidden batch axis
When torch.mean(x) has no dim and the input still has a hidden batch (__batch_index != 233), this warning is followed by emitting reduce_all=1. Reduction does not set support_batch, so the runtime applies it independently to each Mat.n slice; for an [N,C,H,W] input this returns one mean per batch rather than the single mean over all N*C*H*W elements that PyTorch produces.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 37de65b6b2
ℹ️ 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".
| if (r->producer->params.find("__torch_batch_index") != r->producer->params.end()) | ||
| batch_index = r->producer->params.at("__torch_batch_index").i; | ||
|
|
||
| fprintf(pyfp, " ex.input(\"%s\", ncnn.Mat(%s.numpy(), batch_index=%d).clone())\n", input_name.c_str(), input_name.c_str(), batch_index); |
There was a problem hiding this comment.
Avoid passing batch_index for rank-1 inputs
When a torch input is rank 1 but solve_batch_index marks it as batch axis 0 (for example F.embedding/other known batch-first ops do not exempt rank-1 inputs), this generated call becomes ncnn.Mat(in0.numpy(), batch_index=0). The new Python constructor rejects batched conversion for ndim < 2, so the generated *_ncnn.py test/inference script raises before running; the old squeeze(0) path did not fail for a non-size-1 rank-1 tensor. Please keep using 233 or otherwise special-case rank-1 unbatched inputs here.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8481f13748
ℹ️ 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".
| if (axis == batch_index) | ||
| { | ||
| int input_rank = in->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "roll along batch axis %d is not supported\n", batch_index); | ||
| } |
There was a problem hiding this comment.
Reject rolling the hidden batch axis
When torch.roll targets the hidden batch dimension, for example dims=0 on an input whose __batch_index is 0, this only prints a warning and then continues to emit the Slice/Concat roll using the same axis. In the new multi-batch layout that axis is no longer present in the per-sample ncnn Mat, so the converted model rolls a channel/spatial axis (or no equivalent axis) instead of reordering batches; this case should fall back or materialize the batch axis rather than lowering it.
Useful? React with 👍 / 👎.
| if (batch_index != 233) | ||
| fprintf(stderr, "norm along batch axis is not supported yet\n"); | ||
|
|
||
| op->params["1"] = 1; |
There was a problem hiding this comment.
Do not reduce-all norm on hidden batches
For torch.norm(x) with no dim and a hidden batch axis, this warning is followed by reduce_all=1. Because the Reduction layer is still applied per Mat.n batch, the converted model returns one norm per batch instead of the single norm over all elements that PyTorch computes; this should fall back or first materialize the batch dimension.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: abf8599b34
ℹ️ 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".
| if (axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "softmax along batch axis %d is not supported\n", batch_index); | ||
| return; |
There was a problem hiding this comment.
Reject F.softmax over hidden batch axis
When F.softmax is called on the hidden batch dimension, e.g. F.softmax(x, dim=0) after a batch-aware layer on an [N,C,H,W] tensor, this path only logs and returns from write(). The graph has already been rewritten to a new Softmax op, so leaving params unset makes Softmax::load_param use its default axis 0, and because Softmax is still run per Mat.n batch it normalizes a per-sample axis instead of across N. This case needs to fall back/materialize the batch axis rather than emitting a default softmax.
Useful? React with 👍 / 👎.
| if (axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "softmax along batch axis %d is not supported\n", batch_index); | ||
| return; |
There was a problem hiding this comment.
Reject nn.Softmax over hidden batch axis
For nn.Softmax(dim=0) on a tensor whose batch is stored in Mat.n, this branch only prints a warning and returns from write(). Since the rewrite has already created a Softmax op with no axis params, ncnn defaults to axis 0 and applies it independently to each batch slice, so the converted model never computes the PyTorch softmax across the batch dimension.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 47720ecc1c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Line 952 in 51dce2a
When Vulkan executes the new batched fallback for any layer with support_batch == false, top_batch_slot is only a sub-buffer view into top_batch, but VkCompute::record_clone(const VkMat&, VkMat&, ...) recreates its dst with create_like before copying. The copy therefore goes into a fresh allocation instead of the batch slot, so blob_mats_gpu[top_blob_index] = top_batch exposes unwritten data; the same pattern is used in the multi-bottom fallback below.
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a5fb37f86f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "cat along batch axis %d is not supported\n", batch_index); | ||
| continue; |
There was a problem hiding this comment.
Do not keep Concat for batch-axis cat
When torch.cat concatenates along the hidden batch axis, the pass only logs and continues after changing the op to Concat. Since ncnn Concat runs on each per-sample Mat and cannot append the Mat.n batches, a model like torch.cat([a, b], dim=0) with batched inputs keeps/concats per-sample axes instead of producing a larger batch; this should fall back or materialize the batch dimension before lowering.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1d22e09959
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "split along batch axis %d is not supported\n", batch_index); | ||
| continue; |
There was a problem hiding this comment.
Avoid emitting Slice for batch-axis split
When torch.split is requested on the hidden batch axis (for example dim=0 on an [N,C,H,W] tensor carried in Mat.n), this continue runs after the op has already been changed to Slice, so params 0/1 are never written and the original split_size_or_sections/dim params remain. The generated Slice cannot split Mat.n batches and may dereference an empty slices array at runtime instead of falling back or materializing the batch axis.
Useful? React with 👍 / 👎.
| if (axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "normalize along batch axis %d is not supported\n", batch_index); | ||
| return; |
There was a problem hiding this comment.
Do not lower normalize over hidden batches
When F.normalize(..., dim=batch_index) targets the hidden batch dimension, returning here leaves the already-rewritten Normalize op without the params/scale attribute this writer normally emits. That is not a fallback: the saved ncnn layer either fails to load with default scale_data_size=0 or normalizes a per-sample axis, so cross-batch normalization is lost for models that normalize along N.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fe530b10f7
ℹ️ 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 | ||
| static const int n = 1; | ||
| #endif // NCNN_BATCH |
There was a problem hiding this comment.
Keep nstep available when batch is disabled
With NCNN_BATCH=OFF, this #else leaves Mat without an nstep member, but the optimized packing implementations still reference top_blob.nstep/bottom_blob.nstep in their 1-D fast paths (checked ARM, MIPS, LoongArch, and RISC-V packing files). Cross-building those backends with the new NCNN_BATCH option disabled will fail to compile; either keep a non-batch nstep compatibility constant/member or guard those remaining users too.
Useful? React with 👍 / 👎.
|
|
||
| if (axis < 0) | ||
| if (batch_index != 233 && batch_in_shape == 0 && axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "select along batch axis %d is not supported\n", batch_index); | ||
| axis = 0; |
There was a problem hiding this comment.
Reject select on the hidden batch axis
When Tensor.select targets the hidden batch dimension, e.g. x.select(0, 1) for a tensor carried in Mat.n, this logs the unsupported case but then rewrites the op to crop axis 0 of each per-sample Mat. That selects channels/spatial data while keeping all batches, instead of selecting batch 1; this path should fall back/materialize the batch axis rather than continuing with axis = 0.
Useful? React with 👍 / 👎.
| if (batch_index != 233 && batch_in_shape == 0 && axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "chunk along batch axis %d is not supported\n", batch_index); | ||
| axis = 0; | ||
| } |
There was a problem hiding this comment.
Reject chunk on the hidden batch axis
When torch.chunk splits the hidden batch dimension, such as torch.chunk(x, 2, dim=0) for a batched tensor stored in Mat.n, this only prints a warning and then emits an ncnn Slice on axis 0. The slice sizes are computed from the original batch dimension but the runtime axis is a per-sample dimension, so the converted model splits channel/spatial data instead of the batch; this should not be lowered without materializing the batch axis.
Useful? React with 👍 / 👎.
| if (batch_index != 233 && batch_in_shape == 0 && axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "tensor_split along batch axis %d is not supported\n", batch_index); | ||
| axis = 0; | ||
| } |
There was a problem hiding this comment.
Reject tensor_split on hidden batches
For torch.tensor_split(..., dim=batch_index) on a tensor whose batch is hidden in Mat.n, this branch logs the unsupported axis but still proceeds with axis = 0. The generated Slice then applies the batch-derived sections/indices to the first per-sample ncnn axis and never splits Mat.n, so outputs contain incorrectly sliced channels/spatial data rather than batch partitions.
Useful? React with 👍 / 👎.
| if (batch_index != 233 && batch_in_shape == 0 && axis == batch_index) | ||
| { | ||
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "unbind along batch axis %d is not supported\n", batch_index); | ||
| axis = 0; | ||
| } |
There was a problem hiding this comment.
Reject unbind on the hidden batch axis
When torch.unbind is requested over the hidden batch dimension, this logs the unsupported case but still emits a Slice on axis 0. Since ncnn layers operate on each Mat.n batch independently, this unbinds the first per-sample axis instead of producing one output per batch element, so models using unbind(x, dim=0) after a batch-aware op are converted incorrectly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 06a3075f81
ℹ️ 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".
| int input_rank = op->inputs[0]->shape.size(); | ||
| axis = input_rank + axis; | ||
| fprintf(stderr, "glu along batch axis %d is not supported\n", batch_index); | ||
| axis = 0; |
There was a problem hiding this comment.
Reject GLU over the hidden batch axis
When F.glu(..., dim=batch_index) sees a tensor whose batch is stored in Mat.n (__ncnn_batch_in_shape == 0), this warning is followed by forcing axis = 0 and still emitting GLU. Because GLU is then run separately for each batch slice, it splits/gates the first per-sample axis instead of the batch dimension, so models gating across N produce different results; this case needs to fall back or materialize the batch axis rather than lower with axis 0.
Useful? React with 👍 / 👎.
| { | ||
| axis = input_rank + 1 + axis; | ||
| fprintf(stderr, "stack along batch axis %d is not supported\n", batch_index); | ||
| axis = 0; |
There was a problem hiding this comment.
Reject stack on the hidden batch axis
For torch.stack at the hidden batch position, e.g. stacking [N,C,...] tensors with dim=0 while N is in Mat.n, this logs the unsupported axis but then lowers to Concat on axis 0. Concat is applied per batch slice and cannot create a new dimension before the hidden batch, so it concatenates a per-sample axis and preserves the original batch count instead of producing [num_inputs,N,C,...]; this should be rejected or materialized first.
Useful? React with 👍 / 👎.
| fprintf(stderr, "repeat tensor along batch index %d is not supported yet!\n", batch_index); | ||
| continue; |
There was a problem hiding this comment.
Reject repeat factors on the hidden batch axis
When Tensor.repeat asks to repeat the hidden batch dimension with a factor greater than 1, this branch only prints a warning and then drops that axis from new_sizes. The emitted Tile therefore repeats only per-sample dimensions and leaves Mat.n unchanged, so inputs like [N,C].repeat(2,1) fail to double the batch as PyTorch does; the converter should fall back/materialize the batch axis instead of continuing.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 14cb064452
ℹ️ 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".
| if (s == -1) | ||
| s = 0; |
There was a problem hiding this comment.
Preserve dynamic batch sizes when materializing batch
When insert_batch_to_dim materializes a hidden batch for a dynamically shaped tensor, replacing every -1 with 0 makes the generated Reshape copy the corresponding per-sample dimension instead of inferring the extent from Mat.n. For example, an input shaped [-1, 3, 224, 224] that must be converted from hidden batch to an explicit first axis is emitted as [0, 3, 224, 224]; at runtime 0 resolves to the input channel count (3), so any batch size other than 3 causes the batch-to-dim reshape to fail or use the wrong leading dimension. The batch-axis -1 needs to remain inferable (or be resolved from the runtime batch count) rather than being converted to 0.
Useful? React with 👍 / 👎.
No description provided.