The Vanishing Dimension: How keep_dims=false Can Break an Inverse Permutation

Series contents · Numerics and Quantization · 阅读中文版

“Everything passes when we keep the dimensions. Everything fails when we drop them.” It sounds like an unusually polite error message: it practically tells you the answer, yet you can still spend a long time staring at the screen.

The mean itself may be correct. The actual problem is that the third axis before reduction no longer occupies a third position afterward, while the layout-restoration code is still using the old map.

1. An Axis Is More Than an Integer: It Refers to a Coordinate Convention

Take the teaching tensor [A,B,C]=[4,6,9] and reduce B. With dimensions retained, the output is [4,1,9]; without them, it is [4,9].

Both outputs contain the same number of elements, but use different coordinate systems. The former still takes three coordinates, with B restricted to 0; the latter takes only two. Confusing them is like putting a three-dimensional address, including a floor number, into a two-dimensional street map and blaming the navigation software for missing the door.

If the hardware requires a last-axis reduction, we can first transform [A,B,C] into [A,C,B]. B moves from position 1 to position 2. Retaining the reduced axis yields [A,C,1]; the inverse permutation restores [A,1,C]. Only then do we remove the unit dimension to obtain the promised interface shape, [A,C].

The crucial property is not one particular permutation, but preserving rank throughout the intermediate steps:

1

logical input → legal permutation → retain reduced dimensions → inverse permutation → restore interface shape

Positions cannot be removed before the inverse permutation: that permutation still describes axis identities at the old rank.

2. Why the Output Type Cannot Simply Become the Intermediate Type

The historical fix changed the reduction’s intermediate type to one constructed by setting the reduction axes of the effective input shape to one, while taking the original output’s element type. It separates two sources:

• The intermediate shape comes from the current coordinate system and reduction axes.
• The numerical representation comes from the original operator’s output contract.

A convenience helper can easily copy both at once. Copying the final output type seems economical, but introduces its lower-rank shape into an intermediate node too early. Copying the entire input type instead may carry the input’s quantization scale into the reduction result.

A suitable type-construction interface should therefore require explicit choices for shape and element type. Layout encodings, memory spaces, and other information also need a defined policy: inherited, remapped, or recomputed. “Same as before” is not precise enough.

In pseudocode:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# order[new_axis] = old_axis
permuted_input = permute(input, order)
permuted_shape = [input_shape[old_axis] for old_axis in order]
mapped_axes = [new_axis for new_axis, old_axis in enumerate(order)
               if old_axis in normalized_axes]

kept_shape = copy(permuted_shape)
for axis in mapped_axes:
    kept_shape[axis] = 1

temporary_type = tensor(kept_shape, original_result_element_type)
reduced = native_reduce(permuted_input, axes=mapped_axes,
                        keep=true, result_type=temporary_type)
restored = permute(reduced, inverse(order))
result = reshape(restored, original_result_shape)

This conceptual model explicitly computes the temporary shape in the permuted coordinate system. If the element type also contains metadata such as a per-axis quantization axis, that metadata must be remapped too; the example does not expand that part. A common second mistake is recognizing the need for keep=true but still modifying the permuted shape using an axis from before the permutation.

3. When Four Dimensions Become Three, Who Updates the Axes?

Consider another teaching input, [1,4,6,9]. A stage adapting it to three-dimensional hardware removes the leading unit dimension. B, originally axis=2, becomes axis=1. Negative axes should likewise be normalized against the original rank before being mapped into the new coordinate system.

If shape normalization subtracts one and operator-specific legalization subtracts one again using four-dimensional rules, the axis overshoots its destination. Historical changes did move axis legalization after shape normalization. The technical purpose is to establish a single responsibility: wherever the shape changes, the axis mapping must be recorded consistently; later stages consume the new convention.

Three distinct situations must not be blurred into one empty list:

1. The original semantics reduce all axes.
2. The original operation reduces only a unit axis that is later removed.
3. The original operation reduces several axes, some removed and others still present.

If the second case becomes an ordinary empty axes list, and empty axes means all axes, “leave the value unchanged” turns into “collapse the entire tensor to one value.” A historical change introduced a marker for removed reduction axes to preserve exactly this distinction. Identity and Requantization explores its quantization implications.

4. A Correct Permutation Must Also Be Supported

Theoretical support for arbitrary permutations does not mean the backend implements them all. Historical diffs show some cyclic permutations being decomposed into available adjacent-axis swaps; the relevant implementations subsequently evolved in other stages. A test form from one period is not an eternal interface for the entire repository.

The general issue is that legalization must make its helper operations legal as well as the original reduction. If fixing an unsupported reduction introduces an unsupported permutation, the failure has merely moved one stop farther down the pipeline.

A clear contract would require every node produced by legalization to belong to a declared supported set, every new node’s type to be independently inferable from its inputs and attributes, and the final type to equal the original result type exactly. The backend may optimize cost afterward, but must not fill semantic gaps by guessing.

5. Follow the Data Volume When Discussing Performance

Before reduction, the tensor contains A×B×C elements; afterward it contains only A×C. The input permutation may therefore cost far more than the output inverse permutation. Counting “one permutation before and one after” hides a B-fold difference in data volume.

If each element occupies b bytes and both permutations require actual reads and writes, a rough traffic estimate is:

1

extra traffic ≈ 2*b*A*B*C + 2*b*A*C

This is not a measurement. Caches, fusion, layout views, and tiling can all change actual traffic. The estimate helps identify worthwhile questions: can the producer emit the legal layout directly? Can the consumer accept the layout before the inverse permutation? If neither is possible, is a strided reduction an alternative?

Even eliminating a reshape may save no data movement if it was already a view. Explain optimization benefits in terms of transfers, instructions, and buffer lifetimes.

6. Write Tests That Refuse to Be Convenient

Symmetric dimensions are this bug’s friends. If A, B, and C are equal, many wrong permutations still have the same type. If every input value is identical, reducing the wrong axis can still produce the same answer.

Use coordinate-encoded inputs, such as x[a,b,c]=100*a+10*b+c, and give the three dimensions different sizes. A wrong axis then leaves a recognizable numerical fingerprint. Choose coefficients according to the coordinate ranges to avoid overlap; this is only a small teaching encoding.

Historical changes added IR tests for reductions that do not retain dimensions. This establishes a clear regression target. When reading such a patch, ask whether the test checks only the final shape or also the relationship between intermediate permutations and the reduction. The former protects the interface; the latter helps locate the faulty step.

The vanishing dimension has not left the computation. It has only disappeared from the data structure. As long as something downstream still refers to it, the compiler owes it a reliable forwarding address.

7. Use an Axis-Identity Table to Prevent a Second Bad Fix

During debugging, name axes symbolically before assigning numbers. Suppose the original axes are A, B, C. After a permutation, the list is C, A, B. It says new position zero comes from old axis C, and new position one from old axis A. To reduce old axis B, find B in the new list: it is now at position two. The inverse permutation restores A, B, C; it does not simply reuse the original permutation array.

This distinguishes two common arrays: “which old position supplies each new position” and “which new position receives each old position.” They are inverses, coinciding only for certain swaps. Tests that only swap two axes may not fail even when the conventions are confused. A three-axis cyclic permutation forces the distinction to become visible.

Keep labeling the unit reduction axis B instead of deleting it from the list. Axis identity then remains visible throughout: C, A, B becomes C, A, B with B of length one; the inverse permutation restores A, B, C; only the external interface finally removes B. This is why retaining dimensions is easy to prove correct, rather than merely a stylistic preference.

Shared inputs deserve attention too. When shape normalization changes a value’s type, every consumer observes it in the new coordinate system. Updating the reduction’s axes while missing another axis-sensitive consumer can make the local reduction test pass while the full graph remains incorrect. A composition test should attach both a reduction and another axis-sensitive node to the same input, verifying that the shared mapping is applied once and consumed by both.

Stage snapshots are useful here: before and after normalization, and before and after legalization, list each value’s shape, axis identities, and output contract. Avoid relying on intuition while browsing one enormous IR log. Four short tables can often show exactly which stage first violates the invariant.

Back to series contents · Previous · Next

• ← Previous Post
• Next Post →