from __future__ import annotations

from typing import Any

from paddle.common_ops_import import dygraph_only as dygraph_only

from ..fluid.data_feeder import check_dtype as check_dtype
from ..fluid.data_feeder import check_type as check_type
from ..fluid.data_feeder import check_variable_and_dtype as check_variable_and_dtype
from ..fluid.framework import in_dygraph_mode as in_dygraph_mode
from ..fluid.layer_helper import LayerHelper as LayerHelper
from .linalg import dot as dot
from .linalg import matmul as matmul
from .linalg import transpose as transpose
from .manipulation import reshape as reshape
from .manipulation import squeeze as squeeze
from .manipulation import unsqueeze as unsqueeze
from .math import multiply as multiply

def parse_op_labels(labelstr: Any, operand: Any): ...
def parse_labels(labelstr: Any, operands: Any): ...
def validate_rhs(rhs: Any, input_labels: Any, n_bcast_dims: Any) -> None: ...
def build_view(in_labels: Any, out_labels: Any): ...
def build_global_view(nop_labels: Any, rhs: Any, n_bcast_dims: Any): ...
def build_global_shape(g_view: Any, g_labels: Any, op_shapes: Any): ...
def has_duplicated_labels(labels: Any): ...
def diagonalize(labels: Any, operand: Any): ...
def plan_reduce(plan: Any, op: Any, reduce_dims: Any, keepdim: Any): ...
def plan_scalar_prod(plan: Any, op1: Any, op2: Any): ...
def plan_matmul(
    plan: Any, g_view: Any, op1: Any, op2: Any, g_supports: Any, g_shape: Any, I: Any, J1: Any, J2: Any, K: Any
) -> None: ...
def plan_summation(
    plan: Any, g_view: Any, op1: Any, op2: Any, g_supports: Any, g_shape: Any, g_count: Any, n_bcast: Any
) -> None: ...
def rearrange(axes: Any): ...
def plan_broadcast(plan: Any, operands: Any, nop_axes: Any): ...

class Plan:
    env: Any = ...
    steps: Any = ...
    def __init__(self) -> None: ...
    def add_step(self, step: Any) -> None: ...
    def get_var(self, varname: Any): ...
    def set_var(self, varname: Any, var: Any) -> None: ...
    def show(self): ...
    def execute(self): ...

def plan_einsum(operands: Any, g_view: Any, g_shape: Any, g_supports: Any, g_count: Any, n_bcast: Any): ...
def preprocess(equation: Any, *operands: Any): ...
def parse_fake_shape(equation: Any, operands: Any, labels: Any): ...
def rhs_inference(lhs: Any): ...
def gen_equation_for_opteinsum(lhs: Any, rhs: Any): ...
def einsum_v2(equation: Any, *operands: Any): ...
def gen_einsum_op(equation: Any, *operands: Any): ...
def einsum(equation: Any, *operands: Any): ...
