Reordering a Tensor using Einsum

This is an example of how to properly use the einsum function to reorder a tensor.

from openmdao.api import Problem
import numpy as np
from omtools.api import Group
import omtools.api as ot


class ExampleReorderTensor(Group):
    def setup(self):

        # Shape of Tensor
        shape3 = (2, 4, 3)
        c = np.arange(24).reshape(shape3)

        # Declaring tensor
        tens = self.declare_input('c', val=c)

        # Transpose of a tensor
        self.register_output('einsum_reorder2',
                             ot.einsum(tens, subscripts='ijk->kji'))


prob = Problem()
prob.model = ExampleReorderTensor()
prob.setup(force_alloc_complex=True)
prob.run_model()

print('c', prob['c'].shape)
print(prob['c'])
print('einsum_reorder2', prob['einsum_reorder2'].shape)
print(prob['einsum_reorder2'])

c (2, 4, 3)
[[[ 0.  1.  2.]
  [ 3.  4.  5.]
  [ 6.  7.  8.]
  [ 9. 10. 11.]]

 [[12. 13. 14.]
  [15. 16. 17.]
  [18. 19. 20.]
  [21. 22. 23.]]]
einsum_reorder2 (3, 4, 2)
[[[ 0. 12.]
  [ 3. 15.]
  [ 6. 18.]
  [ 9. 21.]]

 [[ 1. 13.]
  [ 4. 16.]
  [ 7. 19.]
  [10. 22.]]

 [[ 2. 14.]
  [ 5. 17.]
  [ 8. 20.]
  [11. 23.]]]