import tensorflow as tf
from ..settings import run as run, ztypes as ztypes
from ..util.warnings import warn_advanced_feature as warn_advanced_feature
from _typeshed import Incomplete
from collections.abc import Callable as Callable
from typing import Any

def constant(value, dtype=..., shape=None, name: str = 'Const', verify_shape=None): ...

pi: Incomplete

def to_complex(number, dtype=...): ...
def to_real(x, dtype=...): ...
def abs_square(x): ...
def nth_pow(x, n): ...
def unstack_x(value: Any, num: Any = None, axis: int = -1, always_list: bool = False, name: str = 'unstack_x'): ...
def stack_x(values, axis: int = -1, name: str = 'stack_x'): ...
def convert_to_tensor(value, dtype=None, name=None, preferred_dtype=None): ...
def safe_where(condition: tf.Tensor, func: Callable, safe_func: Callable, values: tf.Tensor, value_safer: Callable = ...) -> tf.Tensor: ...
def run_no_nan(func, x): ...

class DoNotCompile(Exception): ...

class FunctionWrapperRegistry:
    registries: Incomplete
    allow_jit: bool
    DEFAULT_CACHE_SIZE: int
    do_jit_types: Incomplete
    python_func: Incomplete
    wrapped_func: Incomplete
    force_eager: Incomplete
    wraps: Incomplete
    stateless_args: Incomplete
    function_cache: Incomplete
    currently_traced: Incomplete
    cachesize: Incomplete
    keepalive: Incomplete
    def __init__(self, wraps=None, *, stateless_args=None, cachesize=None, keepalive=None, force_eager=None, **kwargs_user) -> None: ...
    @property
    def do_jit(self): ...
    tf_function_kwargs: Incomplete
    def reset(self, **kwargs_user) -> None: ...
    def set_graph_cache_size(self, cachesize: int | None = None): ...
    @property
    def tf_function(self): ...
    def __call__(self, func): ...

def function(func=None, *, stateless_args=None, cachesize=None, **kwargs): ...
def py_function(func, inp, Tout, name=None): ...
