How to partially flatten a structure, retaining some of the nested structure

username: Samuele_Bolotta
cooked: <p>I was wondering if PyTorch has the equivalent of the TensorFlow “flatten_up_to” function. This allows to partially flatten a structure, retaining some of the nested structure. Thanks!</p>
<p>Line 1158:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/tensorflow/tensorflow/blob/3f878cff5b698b82eea85db2b60d65a2e320850e/tensorflow/python/util/nest.py#L999">
  <header class="source">

      <a href="https://github.com/tensorflow/tensorflow/blob/3f878cff5b698b82eea85db2b60d65a2e320850e/tensorflow/python/util/nest.py#L999" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/tensorflow/tensorflow/blob/3f878cff5b698b82eea85db2b60d65a2e320850e/tensorflow/python/util/nest.py#L999" target="_blank" rel="noopener nofollow ugc">tensorflow/tensorflow/blob/3f878cff5b698b82eea85db2b60d65a2e320850e/tensorflow/python/util/nest.py#L999</a></h4>



    <pre class="onebox">      <code class="lang-py">
        <ol class="start lines" start="989" style="counter-reset: li-counter 988 ;">
            <li>    TypeError: If `check_types` is not `False` and the two structures differ in</li>
            <li>      the type of sequence in any of their substructures.</li>
            <li>    ValueError: If no structures are provided.</li>
            <li>  """</li>
            <li>  return map_structure_with_tuple_paths_up_to(structure[0],</li>
            <li>                                              func,</li>
            <li>                                              *structure,</li>
            <li>                                              **kwargs)</li>
            <li>
            </li>
<li>
            </li>
<li class="selected">def _yield_flat_up_to(shallow_tree, input_tree, is_nested_fn, path=()):</li>
            <li>  """Yields (path, value) pairs of input_tree flattened up to shallow_tree.</li>
            <li>
            </li>
<li>  Args:</li>
            <li>    shallow_tree: Nested structure. Traverse no further than its leaf nodes.</li>
            <li>    input_tree: Nested structure. Return the paths and values from this tree.</li>
            <li>      Must have the same upper structure as shallow_tree.</li>
            <li>    is_nested_fn: Function used to test if a value should be treated as a</li>
            <li>      nested structure.</li>
            <li>    path: Tuple. Optional argument, only used when recursing. The path from the</li>
            <li>      root of the original shallow_tree, down to the root of the shallow_tree</li>
        </ol>
      </code>
    </pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: Matias_Vasquez
cooked: <p>Yes, <a href="https://pytorch.org/docs/stable/generated/torch.flatten.html" rel="noopener nofollow ugc">torch.flatten</a> has a <code>start_dim</code> and a <code>end_dim</code> parameter.</p>
<p>Example:</p>
<pre><code class="lang-python">t = torch.rand((3, 4, 5, 6, 7, 8))

print(t.shape)
t = torch.flatten(t, start_dim=2, end_dim=4)
print(t.shape)
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">torch.Size([3, 4, 5, 6, 7, 8])
torch.Size([3, 4, 210, 8])
</code></pre>

username: Samuele_Bolotta
cooked: <p>Hi! Thanks for your answer.</p>
<p>I don’t fully understand, though: in the TensorFlow function, you can pass two tensors and that will flatten the first up to the the second one. So basically you don’t use explicitly the dimensions, you just pass the structures. In your example however, I wouldn’t know how to flatten tensor t to another structure.</p>

username: Matias_Vasquez
cooked: <p>Hi,</p>
<details>
<summary>
Click here for a long and useless explanation</summary>
<p>I’m going to use the first example from this <a href="https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/contrib/framework/nest/flatten_up_to#examples" rel="noopener nofollow ugc">link</a>.</p>
<pre><code class="lang-python">input_tree = [[[2, 2], [3, 3]], [[4, 9], [5, 5]]]
shallow_tree = [[True, True], [False, True]]

flattened_input_tree = flatten_up_to(shallow_tree, input_tree)
flattened_shallow_tree = flatten_up_to(shallow_tree, shallow_tree)

# Output is:
# [[2, 2], [3, 3], [4, 9], [5, 5]]
# [True, True, False, True]
</code></pre>
<p>So this is what we are given in the example.</p>
<p>The <strong>input</strong> looks like this:</p>
<pre><code class="lang-python">         /-2
      /-|
     |   \-2
   /-|
  |  |   /-3
  |   \-|
  |      \-3
--|
  |      /-4
  |   /-|
  |  |   \-9
   \-|
     |   /-5
      \-|
         \-5
</code></pre>
<p>With shape <code>torch.Size([2, 2, 2])</code> (If we convert this to a torch tensor)</p>
<p>The <strong>shallow tree</strong> looks like this:</p>
<pre><code class="lang-python">      /-True
   /-|
  |   \-True
--|
  |   /-False
   \-|
      \-True
</code></pre>
<p>With shape <code>torch.Size([2, 2])</code> (also already converted to torch tensor).</p>
<p>And the <strong>expected tree</strong> after using the <code>flatten_up_to</code> function is:</p>
<pre><code class="lang-python">      /-2
   /-|
  |   \-2
  |
  |   /-3
  |--|
  |   \-3
--|
  |   /-4
  |--|
  |   \-9
  |
  |   /-5
   \-|
      \-5
</code></pre>
<p>With shape <code>torch.Size([4, 2])</code>.<br>
This means we have to flatten from the left to match the depth of the tree.</p>
<p>Original depth: 3<br>
Shallow depth: 2</p>
<p>We want to “get rid” of 1 dimension, so that they match.</p>
<p>We can achieve this by using</p>
<pre><code class="lang-python">t = torch.flatten(input_tree, end_dim=1)
</code></pre>
<p>To make this more generic, we can create our own function to take care of this:</p>
</details>
<pre><code class="lang-python">def flatten_up_to(shallow_tree, input_tree):
    len_i = len(input_tree.shape)
    len_s = len(shallow_tree.shape)
    d = len_i - len_s
    assert d &gt; 0, "Shallow tree must be smaller than the input tree."

    return torch.flatten(input_tree, end_dim=d)
</code></pre>
<p>You can try and use this and see if it works for you. If this is not the behavior you were expecting, please let me know, perhaps I did not understand correctly what you want/need or how this <code>flatten_up_to</code> is supposed to work <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>
<p>Hope this helps!</p>