from __future__ import annotations

import numpy as np

from .array_base import ArrayBase


class Tensor(ArrayBase):
    def __init__(self, data) -> None:
        super().__init__(data)

    # ---- tensor operations ----

    def contract(self, axes) -> Tensor:
        """Contract the tensor along specified axes."""
        return Tensor(np.tensordot(self.data, self.data, axes=axes))

    def transpose(self, axes=None) -> Tensor:
        """Return the transposed tensor."""
        return Tensor(np.transpose(self.data, axes=axes))

    def einsum(self, subscripts, *operands) -> Tensor:
        """Perform Einstein summation on the tensor."""
        return Tensor(np.einsum(subscripts, self.data, *operands))