tensorplay.stax API

Functions 2

is_available

tensorplay.stax.is_available() → bool[source]: Return whether the native Stax extension is loaded.

stax

tensorplay.stax.stax(graph_module: GraphModule, example_inputs: list[Any], *, mode: str | None = None, options: dict[str, Any] | None = None, name: str | None = None, dynamic: bool | None = None, strict_native: bool = False, **kwargs: Any)[source]

Compile one canonical graph and return an executable callable.

example_inputs and backend options are part of the same contract as metadata in the frontend and uses the native graph when its lowering contract is satisfied. strict_native makes a failed lowering a hard compiler error, so a benchmark can never report the Python GraphModule executor as compiled performance.

Functions 2

is_available

stax

Help improve this page

On this page