Lograr multiplicar la matriz de lotes usando tensordot

Estoy tratando de lograr el mismo comportamiento que la multiplicación de matriz paralela np.matmul usando solo tensordot, punto y remodelación, etc.

La biblioteca a la que estoy traduciendo esto no tiene un matmul que admita la multiplicación paralela, solo punto y tensordot.

Además, quiero evitar iterar sobre la primera dimensión, y quiero hacerlo usando un conjunto de multiplicaciones de matriz y remodelación (quiero que se ejecute tanto con BLAS / GPU como tengo grandes cantidades de pequeñas matrices para calcular en paralelo).

Aquí hay un ejemplo:

import numpy as np

angles = np.array([np.pi/4, 2*np.pi/4, 2*np.pi/4])

vectors = np.array([ [1,0],[1,-1],[-1,0]])

s = np.sin(angles)
c = np.cos(angles)

rotations = np.array([[c,s],[-s,c]]).T

print rotations

print vectors

print("Correct: %s" % np.matmul(rotations, vectors.reshape(3,2,1)))

# I want to do this using tensordot/reshaping, i.e just gemm BLAS operations underneath
print("Wrong: %s" % np.tensordot(rotations, vectors, axes=(1,1)))

El resultado de esto es:

Correct: [[[  7.07106781e-01]
  [  7.07106781e-01]]

 [[  1.00000000e+00]
  [  1.00000000e+00]]

 [[ -6.12323400e-17]
  [ -1.00000000e+00]]]


Wrong: [[[  7.07106781e-01   1.11022302e-16  -7.07106781e-01]
  [ -7.07106781e-01  -1.41421356e+00   7.07106781e-01]]

 [[  6.12323400e-17  -1.00000000e+00  -6.12323400e-17]
  [ -1.00000000e+00  -1.00000000e+00   1.00000000e+00]]

 [[  6.12323400e-17  -1.00000000e+00  -6.12323400e-17]
  [ -1.00000000e+00  -1.00000000e+00   1.00000000e+00]]]

¿Hay alguna manera en que pueda modificar la segunda expresión para obtener el mismo resultado que la primera, simplemente usando punto / tensordot?

Creo que es posible, y he vistoalgunos comentarios en linea, pero nunca hay ejemplos

Respuestas a la pregunta(1)

Su respuesta a la pregunta