title: Equivalence of slicing tensor in Pytorch/ATen C++
body: <p>In Python given a 2-D tensor, we can use <code>tensor[:,:2]</code> to slice the a 2x2 matrix of the first two elements in the top left of the matrix, e.g. :</p>
  
  <pre><code>x = torch.tensor([[-1.4673,  0.9980, -2.1427, -1.1798, -0.0646, -0.2635, -2.8930, -0.2563,
            0.4559, -0.7947, -0.4540,  3.3224,  0.2295,  5.5568, -8.0451, -2.4529,
            4.8724,  2.1640,  3.3255,  0.6693, -1.2362,  4.4713, -3.5547, -0.0528,
            0.1031, -1.2472, -1.6014,  1.8134],
          [ 2.1636, -1.1497, -5.0298,  2.8261, -0.5684,  0.6389,  2.9009, -5.1609,
            1.7358, -3.1819, -0.9877,  5.5598,  6.7142,  4.5704, -1.2683, -5.3046,
            3.0454,  3.2757, -3.2541,  3.6619, -3.6391, -0.2002,  5.7175,  5.7130,
            0.6632, -0.0744, -0.3502,  4.8116]])
  
  y, z = x[:,:2].chunk(2,1)
  
  print(y)
  
  print(z)
  </code></pre>
  
  <p>[out]:</p>
  
  <pre><code>tensor([[-1.4673],
          [ 2.1636]])
  tensor([[ 0.9980],
          [-1.1497]])
  </code></pre>
  
  <p><strong>What is right way to do it in C++ for PyTorch's ATen particularly?</strong></p>
  
  <p>For e.g. in the LSTM, there is the gate.chunk(4,1) function at <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/RNN.cpp#L253" rel="noreferrer">https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/RNN.cpp#L253</a></p>
  
  <p>If I want to do a <code>gate[:,:2].chunk(2,1)</code> to extract different parts of the gates, e.g. <code>auto partial_gates = gates[:,:2].chunk(4, 1);</code>, how can it be done?</p>
  
  <pre><code>template &lt;typename cell_params&gt;
  struct LSTMCell : Cell&lt;std::tuple&lt;Tensor, Tensor&gt;, cell_params&gt; {
    using hidden_type = std::tuple&lt;Tensor, Tensor&gt;;
    hidden_type operator()(const Tensor&amp; input, const hidden_type&amp; hidden, const cell_params&amp; params) const override {
      auto hx = std::get&lt;0&gt;(hidden);
      auto cx = std::get&lt;1&gt;(hidden);
  
      if (input.is_cuda()) {
        auto igates = params.matmul_ih(input);
        auto hgates = params.matmul_hh(hx);
        auto result = at::_thnn_fused_lstm_cell(igates, hgates, cx, params.b_ih, params.b_hh);
        // Slice off the workspace argument (it's needed only for AD).
        return std::make_tuple(std::get&lt;0&gt;(result), std::get&lt;1&gt;(result));
      }
  
      auto gates = params.linear_ih(input) + params.linear_hh(hx);
      auto chunked_gates = gates.chunk(4, 1);
  
      auto partial_gates = gates[:,:2].chunk(4, 1);
  
      auto ingate = chunked_gates[0].sigmoid();
      auto forgetgate = chunked_gates[1].sigmoid();
      auto cellgate = chunked_gates[2].tanh();
      auto outgate = chunked_gates[3].sigmoid();
  
      auto cy = (forgetgate * cx) + (ingate * cellgate);
      auto hy = outgate * cy.tanh();
  
      return std::make_tuple(hy, cy);
    }
  };
  </code></pre>