docs: 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)