Skip to content

cheetah.index_tools

DimSpec module-attribute 

DimSpec: TypeAlias = DimName | int

Equation module-attribute 

Equation: TypeAlias = tuple[tuple[DimSpec, ...], tuple[DimSpec, ...]]

DimName dataclass

A handle representing a named dimension within a specific Dims object.

__init__ 

__init__(_id: int, _dims: Dims, _name: Optional[str] = None) -> None

__deepcopy__ 

__deepcopy__(memo)

__str__ 

__str__() -> str

__repr__ 

__repr__() -> str

Scope dataclass

__init__ 

__init__(_id: int, _dims: Dims) -> None

DimsError

Bases: Exception

ScopeInfo dataclass

eqs instance-attribute 

eqs: list[Equation]

children instance-attribute 

children: list[Scope]

name instance-attribute 

name: str

full_name instance-attribute 

full_name: str

parent class-attribute instance-attribute 

parent: Optional[Scope] = None

__init__ 

__init__(eqs: list[Equation], children: list[Scope], name: str, full_name: str, parent: Optional[Scope] = None) -> None

__repr__ 

__repr__()

Dims

root_scope instance-attribute 

root_scope: Scope = Scope(0, self)

scopes instance-attribute 

scopes: dict[Scope, ScopeInfo] = {root_scope: ScopeInfo(eqs=[], children=[], name='root', full_name='root')}

size_cache instance-attribute 

size_cache: Optional[dict[DimName, int]] = None

curr_scope instance-attribute 

curr_scope: Scope = root_scope

__init__ 

__init__()

new_dim 

new_dim(*args: int | str) -> DimName

new_scope 

new_scope(name: Optional[str] = None) -> Scope

Create a new child scope under the current scope.

If no name is provided, a default base name "scope" is used and a unique suffix is added among siblings.

scope 

scope(scope: Scope)

eq 

eq(lhs: tuple[DimSpec, ...] | DimSpec, rhs: tuple[DimSpec, ...] | DimSpec)

eq_idempotent 

eq_idempotent(lhs: tuple[DimSpec, ...] | DimSpec, rhs: tuple[DimSpec, ...] | DimSpec)

new_dim_with_eq 

new_dim_with_eq(name: str, rhs: tuple[DimSpec, ...] | DimSpec)

size 

size(name: DimName, error_on_missing: bool = True, debug: bool = False) -> int

solve 

solve() -> Indices

init 

init() -> Indices

Indices

sizes instance-attribute 

sizes: dict[DimName, int] = sizes

scopes instance-attribute 

scopes: dict[Scope, ScopeInfo] = scopes

root_scope instance-attribute 

root_scope = root_scope

state instance-attribute 

state: _IndexState = _IndexState(curr_scope=root_scope, eq_flows=dict(), indices_partial=dict(), indices_solved=dict(), eqs=list(eqs), eqs_by_var=eqs_by_var)

__init__ 

__init__(sizes: dict[DimName, int], scopes: dict[Scope, ScopeInfo], root_scope: Scope)

why_partial 

why_partial(name: DimName) -> str

why_solved 

why_solved(name: DimName) -> str

set_index 

set_index(name: DimName, expr: Expr) -> None

scope 

scope(scope: Optional[Scope] = None)

loop 

loop(name: DimName, unroll: Optional[int | Literal[True]] = None)

vectorize 

vectorize(name: DimName, size: int)

debug_check 

debug_check(dims_to_check: list[DimName]) -> None

size 

size(name: DimName) -> int

has_dim 

has_dim(dim: DimName) -> bool

__getitem__ 

__getitem__(name: DimName) -> ch.Expr

__repr__ 

__repr__()

ScopeFn

dims instance-attribute 

dims = dims

scope instance-attribute 

scope = new_scope(__name__)

__init__ 

__init__(dims: Dims, *args, **kwargs)

__call__ 

__call__(ix: Indices, *args, **kwargs)

int_root 

int_root(n: int, r: int) -> Optional[int]

scope_fn 

scope_fn(setup_fn, name: Optional[str] = None)