def test_linear_function():
    """
    **Overview**:
        Test linear function for both forward and backward operation.
    """
    # Generate data.
    w = torch.randn(4, 3, requires_grad=True)
    x = torch.randn(1, 3, requires_grad=False)
    b = torch.randn(4, requires_grad=True)

    # Forward computation graph.
    o = torch.sum(x @ w.t() + b)
    # Backward using auto-grad mechanism.
    o.backward()
    # Save gradients for checking correctness.
    w_grad, b_grad = deepcopy(w.grad), deepcopy(b.grad)
    w.grad, x.grad, b.grad = None, None, None

    # Forward using our defined LinearFunction.
    linear_func = LinearFunction()
    o = torch.sum(linear_func.apply(x, w, b))
    # Backward.
    o.backward()

    # Check whether the results are correct.
    assert x.grad is None
    assert torch.sum(torch.abs(w_grad - w.grad)) < 1e-6
    assert torch.sum(torch.abs(b_grad - b.grad)) < 1e-6

    **Overview**:
        Test linear function for both forward and backward operation.
    
chapter3_obs/gradient.py