Catalyst raises `RuntimeError` when calculating a density matrix on a subset of wires

[albi3ro]

Issue description

Description of the issue - include code snippets and screenshots here if relevant. You may use the following template below

• Expected behavior: (What you expect to happen)

I would expect to get the same output as in the non-qjit scenario:

array([[0.5+0.j, 0.5+0.j],
       [0.5+0.j, 0.5+0.j]])

• Actual behavior: (What actually happens)

traceback below.

• Reproduces how often: (What percentage of the time does it reproduce?)

Every time.

• System information: (post the output of import pennylane as qml; qml.about())

master pennylane, recent catalyst (haven't updated in the last few weeks).

Source code and tracebacks

dev = qml.device('lightning.qubit', wires=2)

@qml.qjit
@qml.qnode(dev)
def circuit(phi):
    qml.Hadamard(0)
    qml.IsingXX(phi, wires=(0,1))
    return qml.density_matrix(wires=0)

circuit(0.5)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[7], line 10
      7     qml.IsingXX(phi, wires=(0,1))
      8     return qml.density_matrix(wires=0)
---> 10 circuit(0.5)

File [pennylane/pennylane/logging/decorators.py:61](/pennylane/pennylane/logging/decorators.py#line=60), in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File [catalyst/frontend/catalyst/jit.py:525](catalyst/frontend/catalyst/jit.py#line=524), in QJIT.__call__(self, *args, **kwargs)
    522     dynamic_args = filter_static_args(args, self.compile_options.static_argnums)
    523     args = promote_arguments(self.c_sig, dynamic_args)
--> 525 return self.run(args, kwargs)

File [/catalyst/debug/instruments.py:143](catalyst/frontend/catalyst/debug/instruments.py#line=142), in instrument.<locals>.wrapper(*args, **kwargs)
    140 @functools.wraps(fn)
    141 def wrapper(*args, **kwargs):
    142     if not InstrumentSession.active:
--> 143         return fn(*args, **kwargs)
    145     with ResultReporter(stage_name, has_finegrained) as reporter:
    146         fn_results, wall_time, cpu_time = time_function(fn, args, kwargs)

File [pennylane/pennylane/logging/decorators.py:61](pennylane/pennylane/logging/decorators.py#line=60), in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File [catalyst/frontend/catalyst/jit.py:713](catalyst/frontend/catalyst/jit.py#line=712), in QJIT.run(self, args, kwargs)
    700 @instrument(has_finegrained=True)
    701 @debug_logger
    702 def run(self, args, kwargs):
    703     """Invoke a previously compiled function with the supplied arguments.
    704 
    705     Args:
   (...)
    710         Any: results of the execution arranged into the original function's output PyTrees
    711     """
--> 713     results = self.compiled_function(*args, **kwargs)
    715     # TODO: Move this to the compiled function object.
    716     return tree_unflatten(self.out_treedef, results)

File [catalyst/frontend/catalyst/compiled_functions.py:335](catalyst/frontend/catalyst/compiled_functions.py#line=334), in CompiledFunction.__call__(self, *args, **kwargs)
    331 abi_args, _buffer = self.args_to_memref_descs(self.restype, dynamic_args)
    333 numpy_dict = {nparr.ctypes.data: nparr for nparr in _buffer}
--> 335 result = CompiledFunction._exec(
    336     self.shared_object,
    337     self.restype,
    338     numpy_dict,
    339     *abi_args,
    340 )
    342 return result

File [catalyst/frontend/catalyst/compiled_functions.py:161](/catalyst/frontend/catalyst/compiled_functions.py#line=160), in CompiledFunction._exec(shared_object, has_return, numpy_dict, *args)
    159 with shared_object as lib:
    160     result_desc = type(args[0].contents) if has_return else None
--> 161     retval = wrapper.wrap(lib.function, args, result_desc, lib.mem_transfer, numpy_dict)
    163 return retval

RuntimeError: [[/catalyst/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp](/catalyst/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp)][Line:238][Function:State] Error in Catalyst Runtime: Invalid size for the pre-allocated state vector

Additional information

Any additional information, configuration or data that might be necessary to reproduce the issue.