def backward(ctx, grad_output):
        """
        **Overview**:
            Backward implementation of linear layer.
        """
        # Get saved parameters back.
        input_, weight = ctx.saved_tensors
        # Initialize gradients to be None.
        grad_input, grad_weight, grad_bias = None, None, None
        # Calculate gradient for input: $$ \nabla input = \nabla output \cdot weight $$
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        # Calculate gradient for weight: $$ \nabla weight = \nabla output^T \cdot input $$
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input_)
        # Calculate gradient for bias: $$ \nabla bias = \sum \nabla output $$
        if ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)
        return grad_input, grad_weight, grad_bias

        **Overview**:
            Backward implementation of linear layer.
        
chapter3_obs/gradient.py