Repository: PennyLaneAI/pennylane (python)

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

Pull request description:
Add is_leaf argument to pytrees\pytrees.py with test at tests\pytrees\test_pytrees.py, fixes #6083
**Context:**
Add is_leaf argument to pytrees\pytrees.py with test at tests\pytrees\test_pytrees.py
fixes #6083 

**Description of the Change:**
Added is_leaf argument to the function pytrees\pytrees.py. Should this argument evaluate to True then no fufther flattening will be done. 
Added test to tests\pytrees\test_pytrees.py, with function test_flatten_is_leaf() having two test cases.
Added changelog with file \doc\releases\changelog-0.38.0.md

**Benefits:**
argument added to the function flatten which acts as a control mechanism to designate a node as leaf node even if that node could be further flattened, e.g., the node is a tuple.

**Possible Drawbacks:**

**Related GitHub Issues:**
Fixes #6083 

**Related Shortcut Stories:**
[sc-70903]

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))