api_reference.md

API reference

The sections below list the public Python interfaces of the gpuastar package.

Core package

gpuastar - Main Package

import gpuastar

The top-level package provides version metadata and runtime bootstrap.

Attributes:

• __version__: Package version string
• __author__: Author name

gpuastar.main - Hydra Entry Point

from gpuastar.main import main

Console script entry point for the Hydra CLI. Not typically called directly.

gpuastar.cli - Command Dispatch

from gpuastar.cli import run_config

Typed command dispatch for Hydra-composed GPUAStar settings.

Functions:

run_config(settings: ConfigSettings) -> None

Run one validated GPUAStar command configuration.

Parameters:

• settings: Validated configuration object (SearchAStarConfig, AOTIExportConfig, or CppWorkflowConfig)

gpuastar.backends - Backend Discovery

from gpuastar.backends import get_backend, available_backends, BackendProtocol, BackendUnavailableError

Backend discovery and selection. Backends are found through the gpuastar.backends entry-point group, so separately installed solver packages (such as gpuastar-cpp and gpuastar-cu) register themselves without changes to the base package.

Classes:

BackendProtocol

Protocol implemented by every registered backend descriptor.

Attributes:

• name: Stable backend identifier

Methods:

• available() -> bool: Return whether this backend is usable

Functions:

available_backends() -> dict[str, BackendProtocol]

Discover installed GPUAStar backends.

Returns: Mapping of backend name to backend descriptor.

get_backend(requested: str = "auto") -> BackendProtocol

Select a backend by name.

Parameters:

• requested: Requested backend name ("py", "cpp", "cu", or "auto")

Returns: Available backend descriptor.

Raises: BackendUnavailableError if an explicitly requested backend is unavailable.

gpuastar.config - Configuration Dataclasses

from gpuastar.config import SearchAStarConfig, AOTIExportConfig, CppWorkflowConfig, SearchRuntimeConfig, load_config

Hydra-composed, Pydantic-validated settings.

Classes:

SearchRuntimeConfig

Mutable one-run execution context derived from validated settings.

Attributes:

• state_progress_callback: Optional callback for per-state progress

SearchInputConfig

Dataset, model, environment, and output paths for search. The base package defines no defaults for the first four fields. Installed environment providers, such as the gpuastar[experiments] extra, inject defaults through Hydra composition.

Fields:

• states: Path to the initial states dataset
• model: Neural network checkpoint path
• env: Environment name
• results_dir: Output directory
• start_idx: Starting state index
• max_states: Maximum states to process (-1 for unlimited)

SearchAlgorithmConfig

Algorithmic controls for weighted A* behavior.

Fields:

• batch_size: Nodes expanded per iteration
• weight: Heuristic weight
• instance_batch_size: Number of search instances run in parallel
• hash_capacity_factor: Hash table capacity factor
• capacity_mode: "adaptive" or "fixed"
• hashtable_algorithm: "linear_probing" or "python_dict"
• solution_move_source: "infer" or "stored"
• duplicate_update_mode: "in_place" or "new_slot"
• goal_stop_mode: "stop_before_goal" or "expand_popped_batch"
• expand_push_mode: "single_pass" or "flatten" - (Python CPU A* only) child-materialization mechanism (single_pass (default) matches the paper's reference implementation, flatten is a heavier flatten/unflatten variant producing identical child nodes)
• nnet_batch_size: Batch size for heuristic evaluation
• search_budget: Maximum node expansions per state (0 = unlimited)

SearchBackendConfig

Backend, device, and implementation selection.

Fields:

• language: "py", "cpp", "cu", or "auto"
• gpu_only: Run the entire search loop in device-resident tensors
• device: Device string
• heuristic_device: Heuristic device override

TorchRuntimeConfig

Torch execution, compilation, precision, and memory controls.

Fields:

• compile_all: Enable torch.compile for all targets
• compile: Compile settings for heuristic evaluation
• compile_env: Compile settings for environment transition functions
• compile_astar_closed: Compile settings for closed-set matching
• fsdp: Fully Sharded Data Parallel (FSDP2)
• tensor_parallel: Tensor parallelism
• dtensor: DTensor placement
• cuda_graphs: CUDA graph capture
• amp: Automatic mixed precision
• amp_dtype: AMP dtype ("float16" or "bfloat16")

SearchProfileConfig

Profiler and operation-timing settings.

Fields:

• enabled: Enable profiling
• steps: Number of profiler steps
• operation_breakdown: Per-operation timing
• cuda_sync_timing: CUDA synchronization around measured intervals
• record_shapes: Record tensor shapes
• memory: Memory profiling
• stack: Stack tracing
• flops: FLOP counting

MemoryDTypeConfig

Dtype names used by tensor-backed A* storage.

Fields:

• node_state, open_index, open_cost, closed_key, closed_cost
• parent_index, parent_move, path_cost, heuristic, total_cost

SearchAStarConfig

Validated configuration for the batch weighted A* runner.

AOTIExportConfig

Validated configuration for exporting heuristic models to AOTInductor packages.

CppWorkflowConfig

Validated configuration for C++ build workflows.

Functions:

load_config(cfg: DictConfig) -> ConfigSettings

Load and validate configuration from Hydra.

Environment Abstraction

gpuastar.environments.environment_abstract

from gpuastar.environments.environment_abstract import Environment, State

Abstract state and environment contracts for GPUAStar search domains.

State (Abstract Base Class)

Abstract base class for hashable environment states.

Abstract Methods:

__hash__(self) -> int

Return a hash that uniquely identifies the state.

__eq__(self, other: object) -> bool

Return True when both states describe the same configuration.

Environment (Abstract Base Class)

Base environment API shared by all search domains.

Attributes:

• dtype: NumPy dtype for state representation
• fixed_actions: Whether the environment has fixed discrete actions
• suggested_compile_max_nodes: Optional node-store capacity hint used to size compiled GPU search artifacts

Abstract Methods:

next_state(self, states: list[State], action: int) -> tuple[list[State], list[float]]

Return successor states and transition costs for an action.

Parameters:

• states: State batch to transform
• action: Zero-based action index

Returns: Pair (next_states, transition_costs) aligned with states.

prev_state(self, states: list[State], action: int) -> list[State]

Return predecessor states produced by reversing an action.

generate_goal_states(self, num_states: int) -> list[State]

Generate num_states solved states.

is_solved(self, states: list[State]) -> NDArray[np.bool_]

Return a solved-state mask with shape (batch_size,).

state_to_nnet_input(self, states: list[State]) -> list[NDArray[np.float32]]

Convert states into float32 arrays accepted by the heuristic model.

get_num_moves(self) -> int

Return the number of discrete actions available.

get_nnet_model(self) -> nn.Module

Return the neural network architecture used for heuristic evaluation.

Concrete Methods:

get_solved_state(self) -> State

Return a single solved state.

generate_states(self, num_states: int, backwards_range: tuple[int, int]) -> tuple[list[State], list[int]]

Sample states by scrambling solved states with reverse moves.

expand(self, states: list[State]) -> tuple[list[list[State]], list[NDArray[np.float32]]]

Generate all fixed-action successors and transition costs.

gpuastar.utils.environments.registry - Environment Registry

from gpuastar.utils.environments.registry import register_environment, register_lazy, get_environment

Registration and lookup of environments by name.

Functions:

register_environment(key: str, cls: type[Environment]) -> None

Register an environment class under a name usable as input.env.

register_lazy(key: str, module_name: str, attr: str | None = None) -> None

Register an environment by module path so that the module is imported only when the environment is requested.

get_environment(env_name: str) -> Environment

Instantiate a registered environment by name.

gpuastar.environments.torch_env

from gpuastar.environments.torch_env import make_torch_env

Adapter constructor for tensor-based environments. The TorchEnv protocol is defined in gpuastar.protocols. The bundled adapter implementations live in the gpuastar-experiments package, so make_torch_env requires the gpuastar[experiments] extra for the example domains.

Protocol: TorchEnv

Attributes:

• device: Torch device
• state_dim: State dimension
• num_moves: Number of moves

Methods:

encode_states(self, states: list[State]) -> Tensor

Encode CPU state objects into a tensor on the device.

decode_states(self, states: Tensor) -> list[State]

Decode a tensor batch into CPU state objects.

is_solved(self, states: Tensor) -> Tensor

Return a solved-state mask.

expand(self, states: Tensor) -> tuple[Tensor, Tensor]

Expand states into all children and transition costs.

Returns: (children, costs) where children has shape (batch_size, num_moves, state_dim) and costs has shape (batch_size, num_moves).

nnet_inputs(self, states: Tensor) -> list[Tensor]

Return tensors formatted for the heuristic network.

Functions:

make_torch_env(env: Environment, device: torch.device) -> TorchEnv

Create a TorchEnv adapter for an Environment.

Backends

Python Backend

from gpuastar.astar.backends.python import bwas_python

bwas_python(args, env, states, operation_config=None)

Python implementation of batch weighted A* search.

Parameters:

• args: SearchRuntimeConfig
• env: Environment instance
• states: List of initial states
• operation_config: Optional OperationConfig

Returns: Tuple (solutions, paths, times, nodes_generated)

C++ Backend

from gpuastar.astar.backends.cpp import bwas_cpp, bwas_cpp_extension

bwas_cpp(args, env, states)

Run the C++ solver. This wrapper validates the requested devices, converts states into the solver payload, dispatches to bwas_cpp_extension, and produces the per-state reports.

bwas_cpp_extension(args, env, states)

Call the solver in the gpuastar_cpp extension module directly and return (solutions, paths, times, nodes_generated).

CUDA Backend

from gpuastar.astar.backends.cuda import bwas_cu

bwas_cu(args, env, states)

Run weighted A* search through the gpuastar-cu kernel backend.

Torch Backend (GPU-Resident)

from gpuastar.astar.backends.torch import bwas_torch

bwas_torch(args, env, states)

Run batch weighted A* search with device-resident tensors and TorchDict.

Utilities

gpuastar.utils.aoti.export

from gpuastar.utils.aoti.export import main as aoti_export_main, run_export

AOTInductor model export utilities.

Functions:

run_export(settings: AOTIExportSettings) -> None

Export a heuristic model checkpoint to an AOTInductor package.

gpuastar.utils.models.neural_networks

from gpuastar.utils.models.neural_networks import load_heuristic_fn, load_nnet, get_device

Neural network loading and heuristic function adapters.

Functions:

load_heuristic_fn(checkpoint, device, on_gpu, model_factory, env, clip_zero=True, batch_size=0, compile_cfg=None, aoti_package_path=None, fuse_conv_bn=False, prune_amount=0.0, sparsity_structured=0.0)

Load a heuristic function from a checkpoint.

Returns: TimedNNetHeuristicFn callable.

load_nnet(checkpoint, model, device)

Load network weights from a checkpoint into a model.

get_device(heuristic_device: str) -> tuple[torch.device, list[str], bool]

Resolve a device specification.

gpuastar.utils.torchdict.core

from gpuastar.utils.torchdict.core import TorchDict, KeySpec, ValueSpec

Open-addressing hash table stored in PyTorch tensors on the selected device.

TorchDict

Compact open-addressing dictionary implemented with PyTorch tensors.

Constructor:

TorchDict(
    key_spec: KeySpec,
    value_spec: ValueSpec,
    capacity: int,
    device: torch.device | str | None = None,
    *,
    hashtable_algorithm: str = "python_dict",
    salt: int | None = None,
)

Methods:

__setitem__(self, keys, values) -> None

Insert or update key-value pairs.

__getitem__(self, keys) -> Tensor

Retrieve values for keys.

__contains__(self, keys) -> Tensor

Check key membership.

get(self, keys, default=None) -> Tensor

Get values with a default for missing keys.

pop(self, keys) -> Tensor

Remove and return values for keys.

keys_tensor(self) -> Tensor

Return a tensor of all keys.

values_tensor(self) -> Tensor

Return a tensor of all values.

items_tensors(self) -> tuple[Tensor, Tensor]

Return (keys, values) tensors.

__len__(self) -> int

Return the number of live items.

stats(self) -> TorchDictStats

Return hash table statistics.

update_min_many(self, keys, values) -> Tensor

Fused closed-set update. For existing keys, the stored value becomes the minimum of the old and new values. Returns a boolean mask marking entries accepted as strict path cost improvements.

KeySpec

Describe how keys are encoded in the tensor storage.

Attributes:

• kind: "int64" or "int64_vec"
• length: Vector width for int64_vec keys

ValueSpec

Describe the dtype and per-key shape of the stored values.

Attributes:

• dtype: Tensor dtype
• shape: Per-key value shape tuple

gpuastar.utils.data.state_datasets - State Datasets

from gpuastar.utils.data.state_datasets import load_states, save_states

State dataset I/O utilities.

Functions:

load_states(path: Path) -> list[State]

Load states from a pickle file.

save_states(states: list[State], path: Path) -> None

Save states to a pickle file.

Protocols

gpuastar.protocols

Structural protocols for optional GPUAStar integrations.

HeuristicFunction

Callable returning cost-to-go estimates for a batch of states.

def __call__(self, states: list[State], is_nnet_format: bool = False) -> NDArray[np.float32]

ModelCallable

Protocol for callable models accepting variadic tensor inputs.

def __call__(self, *inputs: Tensor) -> Tensor | tuple[Tensor, ...] | list[Tensor]

CppPayloadState

Protocol for states that provide their C++ solver payload directly.

def cpp_payload(self) -> Sequence[int]

States without colors or tiles sequences implement this method so the C++ backend can serialize them.

TorchEnv

Protocol for tensor-based environment adapters.

TorchDictView

Minimal protocol for objects that support TorchDict-like views.

Type Aliases

gpuastar.types

Shared type aliases and protocols for the search interfaces.

Type Aliases:

• FloatArray: NDArray[np.float32]
• BoolArray: NDArray[np.bool_]
• Node: TypeVar for node types
• OpenSetElem: tuple[float, int, Node]
• NNetBatch: list[FloatArray]
• UpdateResult: tuple[NNetBatch, FloatArray, BoolArray]
• StateBatch: list[State]
• StateOrNNetBatch: StateBatch | NNetBatch
• ClosedState: TorchDict | TensorClosedState
• JSONValue: JSON-serializable value

Protocols:

HeuristicFn

State-batch heuristic callable.

TimedHeuristicFn

State-batch heuristic with forward-only timing.

NNetHeuristicFn

State-or-network-batch heuristic callable.

TorchHeuristicFn

Tensor-batch heuristic callable.

Dataclasses:

TensorClosedState

Tensor-backed closed-set storage.

TypedDict:

SearchResults

Serialized solver results.

Runtime Checkable Protocols:

StateProgressCallback

Callback invoked after each state completes.

def __call__(self, state_index: int, solved: bool, nodes_generated: int, solve_time_seconds: float) -> None

Type Guards:

is_color_state(state: State) -> TypeGuard[ColorState]

Return True if the state provides a colors attribute.

is_tile_state(state: State) -> TypeGuard[TileState]

Return True if the state provides a tiles attribute.