Эффективное последовательное блочно-диагональное матричное умножение с тензорным потоком
У меня есть блочно-диагональные матрицы вроде
Bdiag1 = [[A, 0, 0], [0, B, 0], [0, 0, C]]
который я могу определить в тензорном потоке, используя LinearOperatorBlockDiag, а затем эффективно вычислить матричное произведение двух таких блочно-диагональных матриц Bdiag1*Bdiag2.
Новая матрица сохраняется как alinear_operator_composition, хотя это все же блочная диагональ. Это означает, что любое последующее матричное умножение не будет иметь диагональной структуры.
Есть ли способ заставить тензорный поток сохранять продукт в виде блочно-диагональной матрицы?
Если это невозможно, есть ли способ определить вектор матриц [A,B,C], где мы можем делать такие вещи, как [A,B,C].matmul([D,E,F]) = [A* D,B* E,C* F]?