def join_differentiable(differentiable_xs, non_differentiable_xs):
        """Reconstitute inputs pytree from differentiable/non-d. partitions."""
        differentiable_leaves = fastmath.tree_leaves(differentiable_xs)
        non_differentiable_leaves = fastmath.tree_leaves(non_differentiable_xs)
        leaves = []
        for is_differentiable in fastmath.tree_leaves(inputs_is_differentiable):
          if is_differentiable:
            leaves.append(differentiable_leaves.pop(0))
          else:
            leaves.append(non_differentiable_leaves.pop(0))
        assert not differentiable_leaves
        assert not non_differentiable_leaves
        tree, _ = fastmath.tree_unflatten(leaves, inputs)
        return tree
Reconstitute inputs pytree from differentiable/non-d. partitions.
trax/layers/research/efficient_attention.py