Unable to extract sub term expectation values when using `cudaq.parallel.mpi`

[bharat-thotakura]

Required prerequisites

• [X] Consult the security policy. If reporting a security vulnerability, do not report the bug using this form. Use the process described in the policy to report the issue.
• [X] Make sure you've read the documentation. Your issue may be addressed there.
• [X] Search the issue tracker to verify that this hasn't already been reported. +1 or comment there if it has.
• [ ] If possible, make a PR with a failing test to give us a starting point to work on!

Describe the bug

Extracting the expectation value of each sub-term in a multi-term Hamiltonian yields zero each time when using the MPI execution mode

Steps to reproduce the bug

Running the following snippet using srun:

import cudaq
from cudaq import spin

cudaq.mpi.initialize()
cudaq.set_target("nvidia", option="mqpu")

@cudaq.kernel
def kernel(angle: float):
    qvector = cudaq.qvector(2)
    x(qvector[0])
    ry(angle, qvector[1])
    x.ctrl(qvector[1], qvector[0])

# Example spin Hamiltonian.
hamiltonian = 5.907 - 2.1433 * spin.x(0) * spin.x(1) - 2.1433 * spin.y(
    0) * spin.y(1) + .21829 * spin.z(0) - 6.125 * spin.z(1)

h_terms = [term for term in hamiltonian.distribute_terms(chunk_count=hamiltonian.get_term_count())]
h_keys = [term.to_string().split()[1] for term in h_terms]

exp_val = cudaq.observe(kernel, hamiltonian, 0.57, execution=cudaq.parallel.mpi)

results_dict = {}
for key, spin_op in zip(h_keys, h_terms):
    results_dict[key] = spin_op.get_coefficient()*exp_val.expectation(sub_term=spin_op)
    
total_exp_val = sum(val for val in results_dict.values())

if cudaq.mpi.rank() == 0:
    print("Expectation value (master rank):", exp_val.expectation())
    print(f"Sub-term expectation value sum (master rank): {total_exp_val.real}\n")
    print(f"All sub-term expectations: {results_dict}")

cudaq.mpi.finalize()  

yields the following output:

Expectation value (master rank): -1.7466087705259774
Sub-term expectation value sum (master rank): 0.0

All sub-term expectations: {'IZ': -0j, 'ZI': 0j, 'YY': -0j, 'XX': -0j, 'II': 0j}

Expected behavior

If execution=cudaq.parallel.thread in cudaq.observe(), then we do get the correct output:

Expectation value (master rank): -1.746608770525977
Sub-term expectation value sum (master rank): -1.7466087705259774

All sub-term expectations: {'IZ': (-5.156643788330257-0j), 'ZI': (-0.18377857511095705+0j), 'YY': (-1.1565933135688304-0j), 'XX': (-1.1565933135688304-0j), 'II': (5.907000220052898+0j)}

which is equivalent to running the following:

results_dict = {}
for key, spin_op in zip(h_keys, h_terms):
    results_dict[key] = cudaq.observe(kernel, spin_op, 0.57).expectation()

which also yields:

Expectation value (master rank): -1.746608770525977
Sub-term expectation value sum (master rank): -1.7466087705259774

All sub-term expectations: {'IZ': -5.156643788330257, 'ZI': -0.18377857511095705, 'YY': -1.1565933135688304, 'XX': -1.1565933135688304, 'II': 5.907000220052898}

Is this a regression? If it is, put the last known working version (or commit) here.

Not a regression

Environment

• CUDA Quantum version: The nvcr.io/nvidia/nightly/cuda-quantum:latest image as of this issue post
• Python version: Python 3.10.12
• Operating system: Ubuntu 22.04.4 LTS

Suggestions

No response