[Bug] Mesolve checks for initial state

[chMoussa]

Short description

When I pass a density matrix of shape (4, 4, 1), the initial checks on mesolve transform it into a (16, 16, 1). When passing jump operators (4,4) the multiplication fail. See test below

What is the expected result?

No response

What is the actual result?

No response

Steps/Code to reproduce

@pytest.mark.parametrize("duration", [torch.rand(1), "duration"])
@pytest.mark.parametrize("input_dm", [True,])
def test_timedependent_with_noise(
    tparam: str,
    param_y: float,
    duration: float | str,
    n_steps: int,
    torch_hamiltonian: Callable,
    hamevo_generator: Sequence,
    sin: tuple,
    sq: tuple,
    input_dm: bool,
) -> None:

    psi_start = random_state(2)
    if input_dm:
        psi_start = density_mat(psi_start)
    dur_val = duration if isinstance(duration, torch.Tensor) else torch.rand(1)

    # simulate with time-dependent solver
    t_points = torch.linspace(0, dur_val[0], n_steps)

    list_ops = Depolarizing(0, error_probability=0.1).tensor(2)
    list_ops = [l.squeeze(-1) for l in list_ops]
    solver = SolverType.DP5_ME
    psi_solver = pyq.mesolve(
        torch_hamiltonian, psi_start, list_ops, t_points, solver
    ).states[-1]

    # simulate with HamiltonianEvolution
    embedding = pyq.Embedding(
        tparam_name=tparam,
        var_to_call={sin[0]: sin[1], sq[0]: sq[1]},
    )
    hamiltonian_evolution = pyq.HamiltonianEvolution(
        generator=hamevo_generator,
        time=tparam,
        duration=duration,
        steps=n_steps,
        solver=solver,
        noise_operators=list_ops,
    )
    values = {"y": param_y, "duration": dur_val}
    psi_hamevo = hamiltonian_evolution(
        state=psi_start, values=values, embedding=embedding
    ).reshape(-1, 1)

    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)

Tracebacks (optional)

No response

Environment details (optional)

No response

Would you like to work on this issue?

None

[chMoussa]

Important for #293

[chMoussa]

@vytautas-a Also if we pass a vector of size (4,1) this works, but adding a batchsize for instance (4,3) will throw an error.

[chMoussa]

Same for passing say a density matrix. Shape (4,4) keeps (4,4) and works. Now with batchsize, say (4,4,3) will throw you an error.

[chMoussa]

So it would be better to set the conventions that sesolve will accept (2^n, batch_size) , mesolve (2^n, 2^n, batch_size)