Element wise dot product on block matricies

    +------+------+   +------+------+    +-------+-------+
    |      |      |   |      |      |    |A11.B11|B12.B12|
    | A11  | A12  |   | B11  | B12  |    |       |       |
    +------+------+ o +------+------| =  +-------+-------+
    |      |      |   |      |      |    |A21.B21|A22.B22|
    | A21  | A22  |   | B21  | B22  |    |       |       |
    +-------------+   +------+------+    +-------+-------+

    ..Note
        If any block of either (A or B) are stored as 1d vectors then we assume
        that it denotes a diagonal matrix efficient dot product using numpy
        broadcasting will be used, i.e. A11*B11

        If either (A or B) of the diagonal matrices are stored as vectors then a more
        efficient dot product using numpy broadcasting will be used, i.e. A11*B11

def block_dot(A, B, diagonal=False):
    """
    Element wise dot product on block matricies

    +------+------+   +------+------+    +-------+-------+
    |      |      |   |      |      |    |A11.B11|B12.B12|
    | A11  | A12  |   | B11  | B12  |    |       |       |
    +------+------+ o +------+------| =  +-------+-------+
    |      |      |   |      |      |    |A21.B21|A22.B22|
    | A21  | A22  |   | B21  | B22  |    |       |       |
    +-------------+   +------+------+    +-------+-------+

    ..Note
        If any block of either (A or B) are stored as 1d vectors then we assume
        that it denotes a diagonal matrix efficient dot product using numpy
        broadcasting will be used, i.e. A11*B11

        If either (A or B) of the diagonal matrices are stored as vectors then a more
        efficient dot product using numpy broadcasting will be used, i.e. A11*B11
    """
    #Must have same number of blocks and be a block matrix
    assert A.dtype is np.dtype('object'), "Must be a block matrix"
    assert B.dtype is np.dtype('object'), "Must be a block matrix"
    assert A.shape == B.shape
    def f(C,D):
        """
        C is an element of A, D is the associated element of B
        """
        Cshape = C.shape
        Dshape = D.shape
        if diagonal and (len(Cshape) == 1 or len(Dshape) == 1\
                or C.shape[0] != C.shape[1] or D.shape[0] != D.shape[1]):
            print("Broadcasting, C: {} D:{}".format(C.shape, D.shape))
            return C*D
        else:
            print("Dotting, C: {} C:{}".format(C.shape, D.shape))
            return np.dot(C,D)
    dot = np.vectorize(f, otypes = [np.object])
    return dot(A,B)