from __future__ import annotations

import enum
from typing import Any

class Type(enum.Enum):
    BIJECTION = ...
    INJECTION = ...
    SURJECTION = ...
    OTHER = ...
    @classmethod
    def is_injective(cls, _type: Any): ...

class Transform:
    def __init__(self) -> None: ...
    def __call__(self, input: Any): ...
    def forward(self, x: Any): ...
    def inverse(self, y: Any): ...
    def forward_log_det_jacobian(self, x: Any): ...
    def inverse_log_det_jacobian(self, y: Any): ...
    def forward_shape(self, shape: Any): ...
    def inverse_shape(self, shape: Any): ...

class AbsTransform(Transform): ...

class AffineTransform(Transform):
    def __init__(self, loc: Any, scale: Any) -> None: ...
    @property
    def loc(self): ...
    @property
    def scale(self): ...

class ChainTransform(Transform):
    transforms: Any = ...
    def __init__(self, transforms: Any) -> None: ...

class ExpTransform(Transform):
    def __init__(self) -> None: ...

class IndependentTransform(Transform):
    def __init__(self, base: Any, reinterpreted_batch_rank: Any) -> None: ...

class PowerTransform(Transform):
    def __init__(self, power: Any) -> None: ...
    @property
    def power(self): ...

class ReshapeTransform(Transform):
    def __init__(self, in_event_shape: Any, out_event_shape: Any) -> None: ...
    @property
    def in_event_shape(self): ...
    @property
    def out_event_shape(self): ...

class SigmoidTransform(Transform): ...
class SoftmaxTransform(Transform): ...

class StackTransform(Transform):
    def __init__(self, transforms: Any, axis: int = ...) -> None: ...
    @property
    def transforms(self): ...
    @property
    def axis(self): ...

class StickBreakingTransform(Transform): ...
class TanhTransform(Transform): ...
