Could you change this code to:

        C = A.mean(dim=(-2,), keepdim=True)
        M_strides = [C.stride(i) for i in range(len(C.shape))]
        M_strides[-2] = 0
        M = C.as_strided(A.shape, M_strides)
        return _svd_lowrank(A, q, niter=niter, M=M)

that will avoid creating a possibly large A - C.

Will do that, but I suspect it will just postpone the inevitable - when matmul sees discontiguous arguments, it blow them up. Will have to double check matmul code, but I vaguely remember it from when I looked at it previously.

Do you ever backprop through pca_lowrank? If not, and A is otherwise not needed, A-C can be computed inplace, avoiding large memory allocation.

Will do that, but I suspect it will just postpone the inevitable - when matmul sees discontiguous arguments, it blow them up

I tested the code locally, all pca tests passed ok.

Yeah, sorry, by "blow up" I meant it will materialize expanded contiguous tensors. So yeah, the tests will pass, but memory usage is possibly still suboptimal.

a memory is contiguous, but tensor a is not contiguous. In contiguous tensors, no 2 logically different elements point to the same memory location, whereas in a all elements point to the same memory location.

Right, the whole point of expand-contiguous conversion is to ensure that tensor data is C-contiguous and can be passed to some low-level routine that is not stride-aware, that is, the low-level routine would access the data memory according to the shape information, not strides. That said, matmul low-level routines ought to be stride-aware (thinking of LAPACK) so that the expand-contiguous conversion is unnecessary. So there could be some ideas for matmul optimization here.

I hear what you are saying, however, typically matmul (let's talk about 2d tensors for simplicity) will be dispatched to a blas library that implements gemm API . As you can see here, lda/ldb that are providing stride-awareness are explicitly constrained to be positive, not 0. That said, instead of dispatching calls with stride=0 args to blas, doing pointwise multiplication of arguments followed by a reduction and expanding back to the necessary size is likely to be faster. Double points if multiplication could be done on the fly during reduction, so that A*B (pointwise-multiplied) did not have to be materialized, but unfortunately pytorch cannot do that, you'd need keops for that.

surprised that blas doesn't support zero stride :( Do we need to write specialized kernels for these cases?

The joys of 80's libraries for you :-) The primitive that we need here to avoid materialization is "broadcasted multiplication followed by a reduction" #32591 (comment), einsum people have been talking about it for a long time. It sounds easy but, given how different broadcasting patterns can be and how generally hard it is to write an efficient reduction for various access patterns, turns out to be pretty hard, there are libraries that are doing it to varying degrees of success (Keops and cutensor that I know of)