How to use the lab.B.sum function in lab

def test_sum():
    a = Dense(np.random.randn(10, 20))
    allclose(B.sum(a, axis=0), np.sum(to_np(a), axis=0))

    for x in [Diagonal(np.array([1, 2, 3]), rows=3, cols=5),
              LowRank(left=np.random.randn(5, 3),
                      right=np.random.randn(10, 3),
                      middle=np.random.randn(3, 3))]:
        allclose(B.sum(x), np.sum(to_np(x)))
        allclose(B.sum(x, axis=0), np.sum(to_np(x), axis=0))
        allclose(B.sum(x, axis=1), np.sum(to_np(x), axis=1))
        allclose(B.sum(x, axis=(0, 1)), np.sum(to_np(x), axis=(0, 1)))

def ratio(a, b): return B.sum(a.diag / b.diag)

def _compute_beta_raised(self):
        beta_norm = B.sqrt(B.maximum(B.sum(B.power(self.beta, 2)),
                                     B.cast(B.dtype(self.beta), 1e-30)))
        return B.power(beta_norm, self.alpha)

def sum(a, axis=None): return B.sum(dense(a), axis=axis)

def elwise(self, x, y):
        return B.expand_dims(B.sum(B.multiply(x, y), axis=1), axis=1)

def qf_diag(a, b, c): return B.sum(B.matmul(B.inverse(a), b) * c, axis=0)

def sum(a, axis=None):
    # Efficiently handle a number of common cases.
    if axis is None:
        return B.sum(B.diag(a))
    elif axis is 0:
        return B.concat(B.diag(a),
                        B.zeros(B.dtype(a), B.shape(a)[1] - B.diag_len(a)),
                        axis=0)
    elif axis is 1:
        return B.concat(B.diag(a),
                        B.zeros(B.dtype(a), B.shape(a)[0] - B.diag_len(a)),
                        axis=0)
    else:
        # Fall back to generic implementation.
        return B.sum.invoke(Dense)(a, axis=axis)