def dot_n(x, y):
    '''given two tensors N x I x K and N x K x J return N dot products

    If either x or y is 2-dimensional, broadcast it over all N.

    Dot products are size N x I x J.
    Example:
    x = np.array([[[1,2], [3,4], [5,6]],[[7,8], [9,10],[11,12]]])
    y = np.array([[[1,2,3], [4,5,6]],[[7,8,9],[10,11,12]]])
    print dot_n(x,y)

    array([[[  9,  12,  15],
        [ 19,  26,  33],
        [ 29,  40,  51]],

       [[129, 144, 159],
        [163, 182, 201],
        [197, 220, 243]]])
    '''
    if x.ndim == 2:
        if y.ndim == 2:
            return np.dot(x, y)
        x3 = False
        y3 = True
        nlen = y.shape[0]
    elif y.ndim == 2:
        nlen = x.shape[0]
        x3 = True
        y3 = False
    else:
        assert x.shape[0] == y.shape[0]
        nlen = x.shape[0]
        x3 = True
        y3 = True
    assert x.shape[1+x3] == y.shape[0+y3]
    n, i, j, k = np.mgrid[0:nlen, 0:x.shape[0+x3], 0:y.shape[1+y3],
                          0:y.shape[0+y3]]
    return np.sum((x[n, i, k] if x3 else x[i,k]) *
                  (y[n, k, j] if y3 else y[k,j]), 3)
given two tensors N x I x K and N x K x J return N dot products

    If either x or y is 2-dimensional, broadcast it over all N.

    Dot products are size N x I x J.
    Example:
    x = np.array([[[1,2], [3,4], [5,6]],[[7,8], [9,10],[11,12]]])
    y = np.array([[[1,2,3], [4,5,6]],[[7,8,9],[10,11,12]]])
    print dot_n(x,y)

    array([[[  9,  12,  15],
        [ 19,  26,  33],
        [ 29,  40,  51]],

       [[129, 144, 159],
        [163, 182, 201],
        [197, 220, 243]]])