from __future__ import annotations

from typing import Any, Optional

def vjp(func: Any, xs: Any, v: Optional[Any] = ...): ...
def jvp(func: Any, xs: Any, v: Optional[Any] = ...): ...

class Jacobian:
    def __init__(self, func: Any, xs: Any, is_batched: bool = ...) -> None: ...
    def __getitem__(self, indexes: Any): ...
    @property
    def shape(self): ...

class Hessian:
    symbolic: Any = ...
    def __init__(self, func: Any, xs: Any, is_batched: bool = ...): ...
    def __getitem__(self, indexes: Any): ...
    @property
    def shape(self): ...

class _Jacobian:
    def __init__(self, func: Any, xs: Any) -> None: ...
    @property
    def shape(self) -> None: ...
    def __getitem__(self, indexes: Any): ...

class _JacobianNoBatch(_Jacobian):
    def __init__(self, func: Any, xs: Any) -> None: ...
    @property
    def shape(self): ...

class _JacobianBatchLast(_Jacobian):
    def __init__(self, func: Any, xs: Any) -> None: ...
    @property
    def shape(self): ...

class _JacobianBatchFirst(_Jacobian):
    def __init__(self, func: Any, xs: Any) -> None: ...
    @property
    def shape(self): ...

def jacobian(func: Any, inputs: Any, create_graph: bool = ..., allow_unused: bool = ...): ...
def batch_jacobian(func: Any, inputs: Any, create_graph: bool = ..., allow_unused: bool = ...): ...
def batch_hessian(func: Any, inputs: Any, create_graph: bool = ..., allow_unused: bool = ...): ...
def hessian(func: Any, inputs: Any, create_graph: bool = ..., allow_unused: bool = ...): ...
def vhp(func: Any, inputs: Any, v: Optional[Any] = ..., create_graph: bool = ..., allow_unused: bool = ...): ...
