PyTorch PR #194077: Why an Output View Made nD Linear Backward Copy Data
PyTorch PR #194077 changes the final reshape in the n-dimensional linear fast path from view to _unsafe_view.
The tensor values and storage layout stay the same. The autograd relationship does not.
Before the change, a call such as:
y = torch.nn.functional.linear(x, weight, bias)
y.relu_()could make autograd treat the returned tensor as a differentiable view of an internal addmm result. The in-place activation then had to rebase that view’s history through CopySlices, adding copy work to backward.
The internal addmm result never escaped from linear. There was no separately visible base tensor whose mutation history needed to be preserved.
The patch keeps the reshape but removes that unnecessary differentiable-view relationship.
The fast path flattens and restores the batch dimensions
For a contiguous n-dimensional input with a compatible bias, PyTorch can implement linear through one matrix operation.
A simplified version is:
input:
[d0, d1, ..., dk, in_features]
flatten:
[d0*d1*...*dk, in_features]
addmm:
[d0*d1*...*dk, out_features]
restore:
[d0, d1, ..., dk, out_features]The flattened input is passed to addmm with the transposed weight and bias.
The matrix result is then reshaped back to the original leading dimensions.
The old return statement used:
return result_flattened.view_symint(result_sizes);A normal view does not copy tensor data. It returns another tensor that references the same storage with different shape metadata.
Avoiding a forward copy was already the point of the fast path. The unexpected cost appeared when a later operation modified the returned view in place.
A view carries an aliasing contract
A regular differentiable view tells autograd that two tensor objects represent the same underlying data.
That relationship matters when both objects can be observed.
Consider:
base = some_operation()
view = base.view(new_shape)
view.relu_()The in-place relu_() changes storage also represented by base.
Autograd cannot pretend that only view changed. It must preserve a correct history for the base-and-view relationship, including how the modified region maps back to the original representation.
PyTorch uses machinery including CopySlices and AsStridedBackward to represent that rebased history.
In the old linear path, the same machinery was activated even though base was only the local variable result_flattened.
User code received the reshaped output. It never received the 2D base separately.
The alias existed physically, but the second user-visible identity did not.
In-place ReLU exposed the unnecessary relationship
An out-of-place activation does not create the same problem:
y = F.linear(x, w, b)
z = torch.relu(y)z is a new result, so autograd can follow the ordinary operation graph.
The regression appears with an in-place consumer:
y = F.linear(x, w, b)
y.relu_()Before this PR, the relevant backward graph contained nodes such as:
AddmmBackward0
ViewBackward0
AsStridedBackward0
CopySlices
TBackward0After the change, the graph contains:
AddmmBackward0
ViewBackward0
UnsafeViewBackward0
ReluBackward0
TBackward0One ViewBackward0 remains because flattening the input is still a meaningful view of a user-provided tensor.
The output-side relationship is different. Its base is a temporary created inside linear, so preserving that base as a separately tracked differentiable view adds work without protecting an observable alias.
_unsafe_view still reshapes the tensor
The name can be misleading.
_unsafe_view is not a copying reshape, and it does not disable gradients.
Its native implementation uses the same underlying view construction used for shape changes. Its autograd derivative reshapes the incoming gradient back to the input shape.
The difference is in alias tracking.
PyTorch’s autograd code generation intentionally excludes _unsafe_view from the normal ADInplaceOrView view set. The operator physically aliases its input but is not treated as a differentiable view for the in-place/view history machinery.
That makes it appropriate only when the caller can prove the hidden base will not be observed independently.
The proof in this PR is local:
result_flattened is created inside linear
↓
the reshaped tensor is returned
↓
result_flattened itself does not escape
↓
no external code can hold both identitiesUsing _unsafe_view on an arbitrary user tensor would not have the same justification.
The sparse branch has the same ownership property
The fast path also supports sparse weight layouts.
That branch transposes the problem because sparse addmm expects the sparse matrix on the left:
addmm(bias_transposed, sparse_weight, flattened_input.T).TThe final transposed result is still derived entirely from an addmm output created inside the function.
Neither the intermediate matrix result nor its temporary transposed identity escapes separately.
The patch therefore uses the same _unsafe_view return for both strided and supported sparse weight paths.
The added tests compare values and gradients for COO, CSR, and CSC weights. They also inspect the backward graph and reject CopySlices or AsStrided nodes.
torch.func had a second implementation to fix
PyTorch’s functorch integration contains a separate linear_hack used under torch.func transforms.
That implementation had its own 3D flatten–addmm–restore sequence and ended with a normal view.
Changing only the core Linear.cpp implementation would therefore leave transformed execution with the old behavior.
The PR updates that return to:
return at::_unsafe_view(result, restored_shape);Public review identified that this parallel path initially lacked direct regression coverage.
The final patch adds a torch.func.grad test with a 3D input, bias, and in-place relu_(). A dispatch recorder confirms that _unsafe_view is executed, and the resulting gradients are compared with an equivalent matrix-multiplication expression.
The test does not explicitly cover every transform, such as vmap, but it pins the separate functorch override that could otherwise drift back to view.
The reported benchmark measured the view cost
A reviewer asked for evidence that removing CopySlices produced a measurable benefit.
The author supplied a CUDA microbenchmark using an RTX 5090 and a large four-dimensional input. It compared three forms:
before:
nD addmm result → view → relu_
workaround:
flatten manually → linear → relu_ → view
after:
nD addmm result → _unsafe_view → relu_The reported results were:
before 16.4 ms 4 aten::copy_ calls
workaround 10.4 ms 0 aten::copy_ calls
after 10.4 ms 0 aten::copy_ callsThe profiler’s summed positive device-allocation counter was also reported as 6000 MiB before and 3000 MiB for the workaround and patched form.
These figures describe one large synthetic workload on one GPU. They are not a general F.linear speedup.
The effect requires the flattening fast path and an in-place consumer that would otherwise trigger view-history rebasing. Different shapes, out-of-place activations, compiled graphs, memory pressure, and backend behavior can change the result.
The benchmark still supports the mechanism: the patched nD path matches the manual workaround and removes the observed copy operations.
The tests protect behavior and graph structure
The landed change includes several kinds of evidence.
Dense tests cover 2D, 3D, 4D, and zero-sized shapes. The 3D and 4D cases directly exercise the changed output reshape and would expose the old CopySlices graph.
The 2D case returns the addmm result directly, so it checks surrounding behavior rather than the modified branch. Zero-sized cases help preserve shape handling but do not provide the same pre-change failure signal.
Sparse tests cover COO, CSR, and CSC weights. They verify values, selected gradients, and the absence of CopySlices and AsStrided in the backward graph.
The functorch test exercises the parallel linear_hack.
Profiler, AOT autograd, and distributed debug expectations are updated from aten::view to aten::_unsafe_view, confirming that the operator sequence is intentionally visible to tracing and dispatch tools.
The operator trace is an observable compatibility boundary
The public Python result does not gain a new alias, shape, or numerical value.
However, code that observes ATen operations can now see:
aten._unsafe_viewwhere it previously saw:
aten.viewPyTorch’s in-tree consumers already handle _unsafe_view, and matmul already uses the operator for internal batch-folding reshapes.
Out-of-tree __torch_dispatch__ subclasses or custom tooling that special-case view without supporting _unsafe_view may still observe a change.
That is a bounded compatibility consideration rather than evidence that the patch is incorrect. The implementation needs the different autograd alias contract to remove CopySlices.
It is also why _unsafe_view should remain an internal, proof-carrying choice rather than a general recommendation to application authors.
What the patch actually establishes
The landed code and tests support the following conclusion:
the output reshape shared storage with a local addmm result
↓
regular view tracking treated both identities as observable
↓
an in-place activation rebased history through CopySlices
↓
the hidden base had no external owner
↓
_unsafe_view preserved reshape gradients without that view historyThe patch does not make every linear call faster.
It does not change the 2D direct-addmm path, and it does not benefit an out-of-place consumer through the same mechanism.
Its contract is narrower:
Track an alias only when another observable tensor identity makes that relationship meaningful.
PyTorch PR #194077 applies that rule to the internal output reshape of nD linear, removing copy work that existed solely to maintain the history of a temporary the caller could never access.
Sources
https://github.com/pytorch/pytorch/pull/194077
https://github.com/pytorch/pytorch/commit/b7354c3f61c5e0d880f1b6d5b887c4f9ad12579f
https://github.com/pytorch/pytorch/blob/b7354c3f61c5e0d880f1b6d5b887c4f9ad12579f/test/test_nn.py