Repository: PennyLaneAI/pennylane

Issue:
Add `is_leaf`optional keyword argument to `qml.pytrees.flatten`
PennyLane has its own tools for working with pytrees similar to jax's [`tree`](https://jax.readthedocs.io/en/latest/jax.tree.html#module-jax.tree) module.  These tools can be used to manipulate data structures.

```pycon
>>> x = {"a": (1, 2), "b": 3}
>>> data, struct = qml.pytrees.flatten(x)
>>> data
[1, 2, 3]
>>> struct
PyTreeStructure(dict, ('a', 'b'), [PyTreeStructure(tuple, None, [PyTreeStructure(), PyTreeStructure()]), PyTreeStructure()])
>>> qml.pytrees.unflatten([4,5,6], struct)
{'a': (4, 5), 'b': 6}
```

These utilities are located in [`pennylane/pytrees/pytrees.py/`](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/pytrees/pytrees.py)

This task is to add a `is_leaf: Optional[Callable[[Any], bool]] = None` similar to that in [`jax.tree.flatten`](https://jax.readthedocs.io/en/latest/_autosummary/jax.tree.flatten.html#jax.tree.flatten).

If provided and returning `True`, this callable should take precedence over whether or not the type has a `flatten_fn` for determining if the object is a terminal leaf.

For example,
```
>>> def is_leaf(x: Any) -> bool:
...     return isinstance(x, tuple)
>>> data, struct = qml.pytrees.flatten(x, is_leaf=is_leaf)
>>> data
[(1, 2), 3]
>>> struct
PyTreeStructure(dict, ('a', 'b'), [PyTreeStructure(), PyTreeStructure()])
```

Note that `(1,2)` is now considered a terminal leaf.

A PR should:

* Tests in [`tests/pytrees/test_pytrees.py`](https://github.com/PennyLaneAI/pennylane/blob/master/tests/pytrees/test_pytrees.py)
*  A changelog entry
* Obey black, pylint, and isort

Discussion:
tonmoy-b: hi @albi3ro and @mudit2812, thanks for the review and suggestions. I have accepted the suggested changes. Do let me know if it would be nice for me to look work on any other contributions, I'd love to get more familiar with the code. 
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6107?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.65%. Comparing base [(`9b403fa`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9b403fa844c5665b6cc991695306d8138c594afe?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`ff55b62`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/ff55b62e7d042af62abd30c878d4762759ff722f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 10 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6107      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         431      432       +1     
  Lines       41889    41697     -192     
==========================================
- Hits        41748    41555     -193     
- Misses        141      142       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6107?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

Alex-Preciado: Thank you so much for this PR @tonmoy-b 🚀 !! We will continue with the review next week. Have a great weekend!!
tonmoy-b: > Thanks for opening this PR @tonmoy-b ! The implementation looks great. I left some suggestions for the changelog. Do you mind also adding a couple of more tests:
> 
> * A test to verify that the pytree structure returned by `flatten` is correct when given an `is_leaf` argument.
> * A test to verify that flattening and then unflattening the same pytree will be equivalent to the original object when `flatten` is given an `is_leaf` argument.
> 
> Other than the above comments, the PR is in great shape and I'm happy to approve once the above comments are resolved 🎉

Hi @mudit2812, @albi3ro  @Alex-Preciado , 

      I have added the mentioned tests into tests\pytrees\test_pytrees.py, as three functions (test_flatten_is_leaf, test_unflatten_is_leaf, test_unflatten_for_structure_is_leaf), kindly take a look. 

Regards



Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e120047dd..dd4421616 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -236,6 +236,9 @@
 * Observable validation for `default.qubit` is now based on execution mode (analytic vs. finite shots) and measurement type (sample measurement vs. state measurement).
   [(#5890)](https://github.com/PennyLaneAI/pennylane/pull/5890)
 
+* Added `is_leaf` parameter to function `flatten` in the `qml.pytrees` module. This is to allow node flattening to be stopped for any node where the `is_leaf` optional argument evaluates to being `True`.
+  [(#6107)](https://github.com/PennyLaneAI/pennylane/issues/6107)
+
 <h3>Breaking changes 💔</h3>
 
 * `GlobalPhase` is considered non-differentiable with tape transforms.
@@ -382,6 +385,7 @@ This release contains contributions from (in alphabetical order):
 Tarun Kumar Allamsetty,
 Guillermo Alonso,
 Utkarsh Azad,
+Tonmoy T. Bhattacharya,
 Gabriel Bottrill,
 Ahmed Darwish,
 Astral Cai,
diff --git a/pennylane/pytrees/pytrees.py b/pennylane/pytrees/pytrees.py
index 6e1581fc8..306c17a00 100644
--- a/pennylane/pytrees/pytrees.py
+++ b/pennylane/pytrees/pytrees.py
@@ -217,11 +217,18 @@ class PyTreeStructure:
 leaf = PyTreeStructure(None, (), [])
 
 
-def flatten(obj: Any) -> tuple[list[Any], PyTreeStructure]:
+def flatten(
+    obj: Any, is_leaf: Optional[Callable[[Any], bool]] = None
+) -> tuple[list[Any], PyTreeStructure]:
     """Flattens a pytree into leaves and a structure.
 
     Args:
         obj (Any): any object
+        is_leaf (Callable[[Any], bool] | None = None): an optionally specified
+            function that will be called at each flattening step. It should return
+            a boolean, with ``True`` stopping the traversal and the whole subtree being
+            treated as a leaf, and ``False`` indicating the flattening should traverse
+            the current object.
 
     Returns:
         List[Any], Union[Structure, Leaf]: a list of leaves and a structure representing the object
@@ -239,14 +246,16 @@ def flatten(obj: Any) -> tuple[list[Any], PyTreeStructure]:
     <PyTree(AdjointOperation, (), (<PyTree(Rot, (Wires([0]), ()), (Leaf, Leaf, Leaf))>,))>
     """
     flatten_fn = flatten_registrations.get(type(obj), None)
-    if flatten_fn is None:
+    # set the flag is_leaf_node if is_leaf argument is provided and returns true
+    is_leaf_node = is_leaf(obj) if is_leaf is not None else False
+    if flatten_fn is None or is_leaf_node:
         return [obj], leaf
     leaves, metadata = flatten_fn(obj)
 
     flattened_leaves = []
     child_structures = []
     for l in leaves:
-        child_leaves, child_structure = flatten(l)
+        child_leaves, child_structure = flatten(l, is_leaf)
         flattened_leaves += child_leaves
         child_structures.append(child_structure)
 


Test patch:
diff --git a/tests/pytrees/test_pytrees.py b/tests/pytrees/test_pytrees.py
index 64b188448..9de8c7f91 100644
--- a/tests/pytrees/test_pytrees.py
+++ b/tests/pytrees/test_pytrees.py
@@ -183,3 +183,82 @@ def test_get_typename_type_invalid():
         ValueError, match=re.escape("'not.a.typename' is not the name of a Pytree type.")
     ):
         get_typename_type("not.a.typename")
+
+
+def test_flatten_is_leaf():
+    """Tests for flatten function's is_leaf parameter"""
+    # case 1 - is_leaf is lambda to mark elements from a list to not flatten
+    item_leaves = [(4, 5), 6, {"a": 1, "b": 2}]
+    z = lambda a: (a in item_leaves)
+    items = [1, 2, 3, (4, 5), 6, 7, {"a": 1, "b": 2}, {"a": 3, "b": 4}]
+    data, _ = flatten(items, z)
+    assert data == [1, 2, 3, (4, 5), 6, 7, {"a": 1, "b": 2}, 3, 4]
+
+    # case 2 - is_leaf is lambda to mark any tuple to not be flatten-ed
+    z = lambda a: isinstance(a, tuple)
+    items = [1, 2, 3, (4, 5), 6, 7, (1, 2, 3, 4), {"a": 1, "b": 2}, {"a": 3, "b": 4}]
+    data, _ = flatten(items, z)
+    assert data == [1, 2, 3, (4, 5), 6, 7, (1, 2, 3, 4), 1, 2, 3, 4]
+
+
+def test_unflatten_is_leaf():
+    """Tests for unflatten function following flatten function's is_leaf parameter
+    objective is to confirm that data processed with the flatten function
+    becomes correctly reconstituted with the unflatten function.
+    """
+    # case 1 - is_leaf given a None argument explcitly
+    z = None
+    items = [1, 2, 3, (4, 5), 6, 7, (1, 2, 3, 4), {"a": 1, "b": 2}, {"a": 3, "b": 4}]
+    data, _structure = flatten(items, z)
+    assert data == [1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 1, 2, 3, 4]
+    reconstituted_data = unflatten(data, _structure)
+    assert reconstituted_data == items
+
+    # case 2 - no is_leaf given
+    items = [1, 2, 3, (4, 5), 6, 7, {"a": 1, "b": 2}, {"a": 3, "b": 4}]
+    data, _structure = flatten(items)
+    assert data == [1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4]
+    reconstituted_data = unflatten(data, _structure)
+    assert reconstituted_data == items
+
+    # case 3 - is_leaf is lambda to mark any tuple to not be flatten-ed
+    z = lambda a: isinstance(a, tuple)
+    items = [1, 2, 3, (4, 5), 6, 7, (1, 2, 3, 4), {"a": 1, "b": 2}, {"a": 3, "b": 4}]
+    data, _structure = flatten(items, z)
+    assert data == [1, 2, 3, (4, 5), 6, 7, (1, 2, 3, 4), 1, 2, 3, 4]
+    reconstituted_data = unflatten(data, _structure)
+    assert reconstituted_data == items
+
+    # case 4 - is_leaf is lambda to mark elements from a list to not flatten
+    item_leaves = [(4, 5), 6, {"a": 1, "b": 2}]
+    z = lambda a: (a in item_leaves)
+    items = [1, 2, 3, (4, 5), 6, 7, {"a": 1, "b": 2}, {"a": 3, "b": 4}]
+    data, _structure = flatten(items, z)
+    assert data == [1, 2, 3, (4, 5), 6, 7, {"a": 1, "b": 2}, 3, 4]
+
+
+def test_unflatten_for_structure_is_leaf():
+    """Tests for unflatten function's structure reconstitution following flatten
+    function's is_leaf parameter. Objective is to confirm that data processed
+    with the flatten function becomes correctly reconstituted with respect to
+    structure with the unflatten function.
+    """
+    # case 1 - no is_leaf given
+    items = (1, 2, (3, 4))
+    data, _structure = flatten(items)
+    assert data == [1, 2, 3, 4]
+    assert _structure == PyTreeStructure(
+        tuple, None, [leaf, leaf, PyTreeStructure(tuple, None, [leaf, leaf])]
+    )
+    reconstituted_data = unflatten([5, 6, 7, 8], _structure)
+    assert reconstituted_data == (5, 6, (7, 8))
+
+    # case 2 - is_leaf is lambda to mark elements from a list to not flatten
+    item_leaves = [(3, 4), 6, {"a": 1, "b": 2}]
+    z = lambda a: (a in item_leaves)
+    items = (1, 2, (3, 4))
+    data, _structure = flatten(items, z)
+    assert data == [1, 2, (3, 4)]
+    assert _structure == PyTreeStructure(tuple, None, [leaf, leaf, leaf])
+    reconstituted_data = unflatten([5, 6, (7, 8)], _structure)
+    assert reconstituted_data == (5, 6, (7, 8))