def test_gae() -> None:
    """
    **Overview**:
        Test the GAE function with randomly generated data.
    """
    # Generate random data with trajectory length 10 and batch size 5.
    T, B = 10, 5
    value = torch.randn(T, B)
    next_value = torch.randn(T, B)
    reward = torch.randn(T, B)
    done = torch.randint(0, 2, (T, B)).to(torch.bool)  # Generate random boolean tensor for done flags.
    traj_flag = torch.randint(0, 2, (T, B)).to(torch.bool)  # Generate random boolean tensor for trajectory flags.
    data = (value, next_value, reward, done, traj_flag)

    # Calculate GAE values.
    gae_value = gae(data)

    # Assert that the calculated GAE values have the correct shape.
    assert gae_value.shape == (T, B)

    **Overview**:
        Test the GAE function with randomly generated data.
    
chapter7_tricks/gae.py