def matrix_trace_impl(a, offset=0):
    """
    Computes the trace of an array.
    """

    _check_linalg_matrix(a, "trace", la_prefix=False)

    if not isinstance(offset, (int, types.Integer)):
        raise NumbaTypeError("integer argument expected, got %s" % offset)

    def matrix_trace_impl(a, offset=0):
        rows, cols = a.shape
        k = offset
        if k < 0:
            rows = rows + k
        if k > 0:
            cols = cols - k
        n = max(min(rows, cols), 0)
        ret = 0
        if k >= 0:
            for i in range(n):
                ret += a[i, k + i]
        else:
            for i in range(n):
                ret += a[i - k, i]
        return ret

    return matrix_trace_impl

    Computes the trace of an array.
    
numba/np/linalg.py