from .array_base import ArrayBase
import numpy as np

class Tensor(ArrayBase):

    def __init__(self, data):
        super().__init__(data)
    
    def contract(self, axes):
        """Contract the tensor along specified axes."""
        return Tensor(np.tensordot(self.data, self.data, axes=axes))
    
    def transpose(self, axes=None):
        """Return the transposed tensor."""
        return Tensor(np.transpose(self.data, axes=axes))
    
    def einsum(self, subscripts, *operands):
        """Perform Einstein summation on the tensor."""
        return Tensor(np.einsum(subscripts, self.data, *operands))