[BUG] Why is `sample_log_prob` of shape `td.shape` unlike the other keys which are `[*td.shape, 1]`?

Describe the bug

Building a ProbabilisticActor module with return_log_prob=True makes it write the log-probability of the distribution sample in a given tensordict. However, unlike the other keys in the tensordict (e.g. done, advantage, value_target, ...) which have shape [*tensordict.shape, 1], it has shape tensordict.shape, i.e., without the additional last dimension.

This creates problems when keys are multiplied for example to compute a simple REINFORCE loss.
Say you have a batch of dimension torch.Size([n]), then tensordict["sample_log_prob"] * tensordict["advantage"] will be a tensor of shape torch.Size([n, n]) and taking the mean over it will be different from what one expects.

I also expect this inconsistency to cause other issues in the future.

To Reproduce

Running

import torch
from tensordict.nn import TensorDictModule, make_functional
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import (
    ProbabilisticActor,
    NormalParamWrapper,
    TanhNormal,
)

env = GymEnv("Pendulum-v1")
td = env.rollout(3)
module = NormalParamWrapper(torch.nn.Linear(3, env.action_spec.shape[-1] * 2))
tensordict_module = TensorDictModule(
    module, in_keys=["observation"], out_keys=["loc", "scale"]
)
td_module = ProbabilisticActor(
    module=tensordict_module,
    spec=env.action_spec,
    in_keys=["loc", "scale"],
    distribution_class=TanhNormal,
    return_log_prob=True,
)
params = make_functional(td_module)
td = td_module(td, params=params)
print(td)

reinforce_loss = td["sample_log_prob"] * td["next", "reward"]
print(f"REINFORCE loss has shape {reinforce_loss.shape}")


print(
    "Mean with wrong shapes",
    (torch.arange(3) * torch.arange(3).unsqueeze(1)).float().mean(),
)
print(
    "Mean with correct shapes",
    (torch.arange(3).unsqueeze(1) * torch.arange(3).unsqueeze(1)).float().mean(),
)

Gives

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        loc: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=cpu,
    is_shared=False)
REINFORCE loss has shape torch.Size([3, 3])
Mean with wrong shapes tensor(1.)
Mean with correct shapes tensor(1.6667)

Expected behavior

sample_log_prob should be written with dimension [*tensordict.shape, 1]

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)