set_composite_lp_aggregate

class tensordict.nn.set_composite_lp_aggregate(mode: bool = True)

Controls whether CompositeDistribution log-probabilities and entropies will be aggregated in a single tensor.

When composite_lp_aggregate() returns True, the log-probs / entropies of CompositeDistribution will be summed into a single tensor with the shape of the root tensordict. This behaviour is being deprecated in favor of non-aggregated log-probs, which offer more flexibility and a somewhat more natural API (tensordict samples, tensordict log-probs, tensordict entropies).

The value of composite_lp_aggregate can also be controlled through the COMPOSITE_LP_AGGREGATE environment variable.

Example

>>> _ = torch.manual_seed(0)
>>> from tensordict import TensorDict
>>> from tensordict.nn import CompositeDistribution, set_composite_lp_aggregate
>>> import torch
>>> from torch import distributions as d
>>> params = TensorDict({
...     "cont": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)},
...     ("nested", "disc"): {"logits": torch.randn(3, 10)}
... }, [3])
>>> dist = CompositeDistribution(params,
...     distribution_map={"cont": d.Normal, ("nested", "disc"): d.Categorical})
>>> sample = dist.sample((4,))
>>> with set_composite_lp_aggregate(False):
...     lp = dist.log_prob(sample)
...     print(lp)
TensorDict(
    fields={
        cont_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                disc_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)
>>> with set_composite_lp_aggregate(True):
...     lp = dist.log_prob(sample)
...     print(lp)
tensor([[-2.0886, -1.2155, -0.0414],
        [-2.8973, -5.5165,  2.4402],
        [-0.2806, -1.2799,  3.1733],
        [-3.0407, -4.3593,  0.5763]])