def test_pg():
    """
    **Overview**:
        Test function of PG, for both forward and backward operations.
    """
    # batch size=4, action=32
    B, N = 4, 32
    # Generate logit, action, return_.
    logit = torch.randn(B, N).requires_grad_(True)
    action = torch.randint(0, N, size=(B, ))
    return_ = torch.randn(B) * 2
    # Compute PG error.
    data = pg_data(logit, action, return_)
    loss = pg_error(data)
    # Assert the loss is differentiable.
    assert all([l.shape == tuple() for l in loss])
    assert logit.grad is None
    total_loss = sum(loss)
    total_loss.backward()
    assert isinstance(logit.grad, torch.Tensor)

    **Overview**:
        Test function of PG, for both forward and backward operations.
    
chapter1_overview/pg.py