You gotta know when to fold ’em
With your diagram in place, choose the shape whose branches you want to add the folding function to. Once you select the shape, you should now see that the format panel on the right of your workspace has now changed to show you style options for your shape. Click on Properties, then scroll down the list until you see Tree Folding, and select that.
Take a look at your shape. It should now have a minimize icon on it that is clickable. When you click, this will collapse or expand all the shapes in the branches under it. Voilá!