import tensorflow as tf

    x_shape = tuple(x_shape)
    y_shape = []
    for i, s in zip(int_shape(y), tf.unstack(tf.shape(y))):
      if i is not None:
        y_shape.append(i)
      else:
        y_shape.append(s)
    y_shape = tuple(y_shape)
    y_permute_dim = list(range(get_ndim(y)))
    y_permute_dim = [y_permute_dim.pop(-2)] + y_permute_dim
    xt = tf.reshape(x, [-1, x_shape[-1]])
    yt = tf.reshape(tf.transpose(y, perm=y_permute_dim), [y_shape[-2], -1])
    return tf.reshape(
        tf.matmul(xt, yt), x_shape[:-1] + y_shape[:-2] + y_shape[-1:])
  out = tf.matmul(x, y)
  return out


def get_ndim(x):
  """Returns the number of axes in a tensor, as an integer.

  Parameters
