For loop on slices of a tensor creates multiple intermediate gradients

Jinghan_Yao: <p>Hi, I am trying to use slicing to reduce the peak memory in my model, a simple example of what I am doing is the following:</p>
<pre><code class="lang-auto"># x: Input, [10000, 256]
# w1: Weight, [256, 1024]
# w2: Weight, [1024, 256]
y = torch.matmul(torch.matmul(x, w1), w2)
</code></pre>
<p>Since directly calculating y will create an intermediate activation of <code>[10000, 1024]</code>, I want to use a for loop and hope it can reduce the peak memory:</p>
<pre><code class="lang-auto">chunk_size = 2500
num_chunks = int(x.shape[0] / chunk_size)
y = torch.empty_like(x)
for i in range(num_chunks):
    st = i * chunk_size
    ed = st + chunk_size
    y[st:ed] = torch.matmul(torch.matmul(x[st:ed], w1), w2)
</code></pre>
<p>By doing so, ideally, each for loop only creates an intermediate activation of <code>[2500, 1024]</code>, therefore the peak memory can be reduced.</p>
<p>However, I use pytorch profiler to profile the memory, I found that despite the forward process seeing a decrease in memory, the backward process somehow gets a new memory spike as shown below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/5/65a9b0f5a8079b60c58bf3989e8317aad3ae7d15.png" data-download-href="https://discuss.pytorch.org/uploads/default/65a9b0f5a8079b60c58bf3989e8317aad3ae7d15" title="Untitled"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/5/65a9b0f5a8079b60c58bf3989e8317aad3ae7d15_2_690x450.png" alt="Untitled" data-base62-sha1="evlMKGrQGfjYInuGRWOn56kJRL7" width="690" height="450" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/5/65a9b0f5a8079b60c58bf3989e8317aad3ae7d15_2_690x450.png, https://discuss.pytorch.org/uploads/default/original/3X/6/5/65a9b0f5a8079b60c58bf3989e8317aad3ae7d15.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/5/65a9b0f5a8079b60c58bf3989e8317aad3ae7d15.png 2x" data-dominant-color="D8C2B6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Untitled</span><span class="informations">989×646 177 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The marked circles are named like <code>wrapper_CompositeExplicitAutograd__slice_backward</code>, the full description is:</p>
<pre><code class="lang-auto">CUDACachingAllocator.cpp:0:c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::malloc(int, unsigned long, CUstream_st*)
:0:c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::malloc(void**, int, unsigned long, CUstream_st*)
:0:c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::allocate(unsigned long) const
:0:at::TensorBase at::detail::_empty_generic&lt;long&gt;(c10::ArrayRef&lt;long&gt;, c10::Allocator*, c10::DispatchKeySet, c10::ScalarType, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::detail::empty_generic(c10::ArrayRef&lt;long&gt;, c10::Allocator*, c10::DispatchKeySet, c10::ScalarType, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::detail::empty_cuda(c10::ArrayRef&lt;long&gt;, c10::ScalarType, std::optional&lt;c10::Device&gt;, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::detail::empty_cuda(c10::ArrayRef&lt;long&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::native::empty_cuda(c10::ArrayRef&lt;long&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
RegisterCUDA.cpp:0:at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA_memory_format_empty(c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
RegisterCUDA.cpp:0:c10::impl::wrap_kernel_functor_unboxed_&lt;c10::impl::detail::WrapFunctionIntoFunctor_&lt;c10::CompileTimeFunctionPointer&lt;at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;), &amp;at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA_memory_format_empty&gt;, at::Tensor, c10::guts::typelist::typelist&lt;c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt; &gt; &gt;, at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)&gt;::call(c10::OperatorKernel*, c10::DispatchKeySet, c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::_ops::empty_memory_format::redispatch(c10::DispatchKeySet, c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
RegisterBackendSelect.cpp:0:c10::impl::wrap_kernel_functor_unboxed_&lt;c10::impl::detail::WrapFunctionIntoFunctor_&lt;c10::CompileTimeFunctionPointer&lt;at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;), &amp;at::(anonymous namespace)::empty_memory_format&gt;, at::Tensor, c10::guts::typelist::typelist&lt;c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt; &gt; &gt;, at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)&gt;::call(c10::OperatorKernel*, c10::DispatchKeySet, c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::_ops::empty_memory_format::call(c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;, std::optional&lt;c10::MemoryFormat&gt;)
??:0:at::native::zeros_symint(c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)
RegisterCompositeExplicitAutograd.cpp:0:c10::impl::wrap_kernel_functor_unboxed_&lt;c10::impl::detail::WrapFunctionIntoFunctor_&lt;c10::CompileTimeFunctionPointer&lt;at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;), &amp;at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeExplicitAutograd__zeros&gt;, at::Tensor, c10::guts::typelist::typelist&lt;c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt; &gt; &gt;, at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)&gt;::call(c10::OperatorKernel*, c10::DispatchKeySet, c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)
??:0:at::_ops::zeros::redispatch(c10::DispatchKeySet, c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)
RegisterBackendSelect.cpp:0:c10::impl::wrap_kernel_functor_unboxed_&lt;c10::impl::detail::WrapFunctionIntoFunctor_&lt;c10::CompileTimeFunctionPointer&lt;at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;), &amp;at::(anonymous namespace)::zeros&gt;, at::Tensor, c10::guts::typelist::typelist&lt;c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt; &gt; &gt;, at::Tensor (c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)&gt;::call(c10::OperatorKernel*, c10::DispatchKeySet, c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)
??:0:at::_ops::zeros::call(c10::ArrayRef&lt;c10::SymInt&gt;, std::optional&lt;c10::ScalarType&gt;, std::optional&lt;c10::Layout&gt;, std::optional&lt;c10::Device&gt;, std::optional&lt;bool&gt;)
??:0:at::native::slice_backward(at::Tensor const&amp;, c10::ArrayRef&lt;long&gt;, long, long, long, long)
RegisterCompositeExplicitAutograd.cpp:0:at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeExplicitAutograd__slice_backward(at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)
RegisterCompositeExplicitAutograd.cpp:0:c10::impl::wrap_kernel_functor_unboxed_&lt;c10::impl::detail::WrapFunctionIntoFunctor_&lt;c10::CompileTimeFunctionPointer&lt;at::Tensor (at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt), &amp;at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeExplicitAutograd__slice_backward&gt;, at::Tensor, c10::guts::typelist::typelist&lt;at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt&gt; &gt;, at::Tensor (at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)&gt;::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)
:0:at::Tensor c10::callUnboxedKernelFunction&lt;at::Tensor, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt&gt;(void*, c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;&amp;&amp;, long&amp;&amp;, c10::SymInt&amp;&amp;, c10::SymInt&amp;&amp;, c10::SymInt&amp;&amp;)
:0:at::Tensor c10::Dispatcher::redispatch&lt;at::Tensor, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt&gt;(c10::TypedOperatorHandle&lt;at::Tensor (at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)&gt; const&amp;, c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt) const
??:0:at::_ops::slice_backward::redispatch(c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)
VariableType_3.cpp:0:torch::autograd::VariableType::(anonymous namespace)::slice_backward(c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)
VariableType_3.cpp:0:c10::impl::wrap_kernel_functor_unboxed_&lt;c10::impl::detail::WrapFunctionIntoFunctor_&lt;c10::CompileTimeFunctionPointer&lt;at::Tensor (c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt), &amp;torch::autograd::VariableType::(anonymous namespace)::slice_backward&gt;, at::Tensor, c10::guts::typelist::typelist&lt;c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt&gt; &gt;, at::Tensor (c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)&gt;::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)
:0:at::Tensor c10::callUnboxedKernelFunction&lt;at::Tensor, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt&gt;(void*, c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;&amp;&amp;, long&amp;&amp;, c10::SymInt&amp;&amp;, c10::SymInt&amp;&amp;, c10::SymInt&amp;&amp;)
??:0:at::_ops::slice_backward::call(at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt;, long, c10::SymInt, c10::SymInt, c10::SymInt)
:0:torch::autograd::generated::details::slice_backward_wrapper(at::Tensor const&amp;, c10::ArrayRef&lt;c10::SymInt&gt; const&amp;, long, std::optional&lt;c10::SymInt&gt;, std::optional&lt;c10::SymInt&gt;, c10::SymInt)
??:0:torch::autograd::generated::SliceBackward0::apply(std::vector&lt;at::Tensor, std::allocator&lt;at::Tensor&gt; &gt;&amp;&amp;)
:0:torch::autograd::Node::operator()(std::vector&lt;at::Tensor, std::allocator&lt;at::Tensor&gt; &gt;&amp;&amp;)
??:0:torch::autograd::Engine::evaluate_function(std::shared_ptr&lt;torch::autograd::GraphTask&gt;&amp;, torch::autograd::Node*, torch::autograd::InputBuffer&amp;, std::shared_ptr&lt;torch::autograd::ReadyQueue&gt; const&amp;)
??:0:torch::autograd::Engine::thread_main(std::shared_ptr&lt;torch::autograd::GraphTask&gt; const&amp;)
??:0:torch::autograd::Engine::execute_with_graph_task(std::shared_ptr&lt;torch::autograd::GraphTask&gt; const&amp;, std::shared_ptr&lt;torch::autograd::Node&gt;, torch::autograd::InputBuffer&amp;&amp;)
??:0:torch::autograd::python::PythonEngine::execute_with_graph_task(std::shared_ptr&lt;torch::autograd::GraphTask&gt; const&amp;, std::shared_ptr&lt;torch::autograd::Node&gt;, torch::autograd::InputBuffer&amp;&amp;)
??:0:torch::autograd::Engine::execute(std::vector&lt;torch::autograd::Edge, std::allocator&lt;torch::autograd::Edge&gt; &gt; const&amp;, std::vector&lt;at::Tensor, std::allocator&lt;at::Tensor&gt; &gt; const&amp;, bool, bool, bool, std::vector&lt;torch::autograd::Edge, std::allocator&lt;torch::autograd::Edge&gt; &gt; const&amp;)
??:0:torch::autograd::python::PythonEngine::execute(std::vector&lt;torch::autograd::Edge, std::allocator&lt;torch::autograd::Edge&gt; &gt; const&amp;, std::vector&lt;at::Tensor, std::allocator&lt;at::Tensor&gt; &gt; const&amp;, bool, bool, bool, std::vector&lt;torch::autograd::Edge, std::allocator&lt;torch::autograd::Edge&gt; &gt; const&amp;)
??:0:THPEngine_run_backward(_object*, _object*, _object*)
</code></pre>
<p>And as I use smaller <code>chunk_size</code> (i.e. more slices), this part of memory keeps increasing.</p>
<p>From what I understand, during backpropagation, pytorch is trying to create a gradient for each slice, and store them all in memory before accumulating them. I am wondering why it cannot just perform the exact for loop during backward, and accumulating gradients after each loop, since in this case, the peak memory can be kept low.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/jinghan_yao">@Jinghan_Yao</a>,</p>
<p>If you want to ‘chunk’ the following operation, you could use <code>torch.func.vmap</code> to vectorize over the operation and use the <code>chunk_size</code> arg of <code>torch.func.vmap</code> to chunk your operation.</p>
<aside class="quote no-group" data-username="Jinghan_Yao" data-post="1" data-topic="204442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jinghan_yao/48/70657_2.png" class="avatar"> Jinghan_Yao:</div>
<blockquote>
<pre><code class="lang-auto"># x: Input, [10000, 256]
# w1: Weight, [256, 1024]
# w2: Weight, [1024, 256]
y = torch.matmul(torch.matmul(x, w1), w2)
</code></pre>
</blockquote>
</aside>
<pre><code class="lang-auto">from torch.func import vmap
def f(x,w1,w2):
  return torch.matmul(torch.matmul(x,w1),w2)

y = vmap(f, in_dims=(0,None,None), chunk_size=2500)(x,w1,w2)
</code></pre>
<p>Also, you could try running this operation within a <code>torch.no_grad()</code> to disable gradient calculations.</p>

Jinghan_Yao: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>, thanks. Since I need the gradient, I set the <code>torch.no_grad()</code> to <code>False</code>, however, I got an error when I pass my forward function to <code>vmap</code>:</p>
<pre><code class="lang-auto">RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. For more details, please see https://pytorch.org/docs/master/notes/extending.func.html
</code></pre>
<p>As indicated in the linked document, it seems I need to implement the backward function? I am using torch==2.2.1.</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="Jinghan_Yao" data-post="3" data-topic="204442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jinghan_yao/48/70657_2.png" class="avatar"> Jinghan_Yao:</div>
<blockquote>
<p><code>autograd.Function</code></p>
</blockquote>
</aside>
<p>Are you using a custom <code>torch.autograd.Function</code>? If so, in that case you’ll need to manually define the backward method in order to get it to work with <code>torch.func</code>.</p>

Jinghan_Yao: <p>Oh I see. Yes I use a custom function. Thanks for the instruction!!!</p>

AlphaBetaGamma96: <p>Also, I just had that idea that if you want to only have the intermediate shape of <code>[2500,1024]</code>, couldn’t you just re-define your matmul as an einsum-op?</p>
<aside class="quote no-group" data-username="Jinghan_Yao" data-post="1" data-topic="204442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jinghan_yao/48/70657_2.png" class="avatar"> Jinghan_Yao:</div>
<blockquote>
<p>By doing so, ideally, each for loop only creates an intermediate activation of <code>[2500, 1024]</code>, therefore the peak memory can be reduced.</p>
</blockquote>
</aside>
<p>So,</p>
<aside class="quote no-group" data-username="Jinghan_Yao" data-post="1" data-topic="204442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jinghan_yao/48/70657_2.png" class="avatar"> Jinghan_Yao:</div>
<blockquote>
<p><code>torch.matmul(torch.matmul(x[st:ed], w1), w2)</code></p>
</blockquote>
</aside>
<p>written as,</p>
<pre><code class="lang-auto">torch.einsum("bi,ij,jk-&gt;bk",x,w1,w2)
</code></pre>
<p>and let <code>einsum</code> pick the optional ordering of operations, and just for-loop over it as before?</p>