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: zouzias
cooked: <p>You can take a look at the <code>pytree</code> abstraction that is used internally by PyTorch: <a href="https://github.com/pytorch/pytorch/blob/v1.11.0/torch/utils/_pytree.py#L7" class="inline-onebox" rel="noopener nofollow ugc">pytorch/_pytree.py at v1.11.0 · pytorch/pytorch · GitHub</a></p>
<p>Here is a nice overview of pytrees, if you haven’t seen it before: <a href="https://jax.readthedocs.io/en/latest/pytrees.html" class="inline-onebox" rel="noopener nofollow ugc">Pytrees — JAX documentation</a></p>