def METHOD_NAME(
    distr_out: DistributionOutput,
    data: Tensor,
    loc: List[Union[None, Tensor]],
    scale: List[Union[None, Tensor]],
    expected_batch_shape: Tuple,
    expected_event_shape: Tuple,
):
    args_proj = distr_out.get_args_proj()
    args_proj.initialize()
    args = args_proj(data)
    assert distr_out.event_shape == expected_event_shape
    for l, s in product(loc, scale):
        distr = distr_out.distribution(args, loc=l, scale=s)
        assert distr.batch_shape == expected_batch_shape
        assert distr.event_shape == expected_event_shape
        x = distr.sample()
        assert x.shape == distr.batch_shape + distr.event_shape
        loss = distr.loss(x)
        assert loss.shape == distr.batch_shape
        x1 = distr.sample(num_samples=1)
        assert x1.shape == (1,) + distr.batch_shape + distr.event_shape
        x3 = distr.sample(num_samples=3)
        assert x3.shape == (3,) + distr.batch_shape + distr.event_shape