SGraph

SGraph is a design-stage Rust frontend for OpenXLA. It replaces JAX's frontend layer while retaining StableHLO, XLA, and PJRT as the compiler and runtime.

SGraph constructs complete static, replica-local graphs. Every Graph is the Jaxpr-like function language value; every tensor type has only a local dtype and shape; and all cross-device communication is an explicit collective primitive. There are no global arrays, PartitionSpec, sharding propagation, implicit resharding, eager tensors, graph breaks, or ambient tracing stacks.

Status

The first graph-building, autodiff, StableHLO, and CPU execution slice is implemented. The public API is still experimental and has no compatibility promise yet.

Rust imports keep graph expressions compact:

use sgraph::numpy as np;

let y = np::exp(np::tanh(&x));

Build And Rust Analyzer

SGraph uses Bazel 8.7.0 and rules_rust directly; there is intentionally no Cargo.toml. The workspace requests Rust 1.97.0 with edition 2024 and gives each Rust source file its own Bazel library and unit-test target.

Build the facade and run all tests with:

bazel test //...

Format the repository with:

bazel run //:format

This command uses rules_rust's pinned rustfmt; it does not use a Rust tool from the user's environment.

Set up the Bazel-aware Rust Analyzer launchers for VS Code with:

bazel run @rules_rust//tools/rust_analyzer:setup -- vscode

Run setup after cloning, changing the Rust toolchain, or running bazel clean --expunge, then reload the VS Code window. The generated launchers keep Rust Analyzer's project discovery, diagnostics, formatting, and proc-macro support synchronized with Bazel.

Fixed Decisions

• Rust is the frontend language.
• SGraph replaces only the JAX frontend. StableHLO remains the interchange, XLA remains the optimizer/code generator, and PJRT remains the runtime API.
• Graph is a typed Jaxpr-like function. Nested computation is carried by graph-bearing primitive parameters such as call_graph, cond_graph, and body_graph.
• All source functions are constructed in graph mode and lower directly to StableHLO; there is no frontend eager mode or host graph evaluator.
• Primitive abstract evaluation, AD rules, and lowering rules use explicit registries; graph transforms operate directly on completed graphs.
• jvp, linearize, transpose, and vjp are direct graph-to-graph transformations. Lowering walks the completed graph directly.
• Rust operators and namespace functions insert checked graph primitives through Tracer values; invalid graph construction raises immediately, and no numeric computation executes in the frontend.
• Every graph is per-device code. A mesh configures XLA replicas, not frontend sharding; collectives are the only communication mechanism.
• Bazel is the sole build system. Rust uses direct, Bazel-native rules_rust targets in non-Cargo dependency mode: every Rust source file has its own library and unit-test target, while src/lib.rs is a facade over explicit file-target dependencies. There are no Cargo manifests, crate_universe, Cargo build scripts, or Cargo-generated dependency graph.

Design Index

• Frontend core: historical core-design proposal. The current graph and autodiff architecture is described by the architecture overview and autodiff notes.
• System design: project boundary and end-to-end ownership.
• Architecture overview: short component map.
• Autodiff: current JVP/VJP graph-transform design.
• StableHLO, XLA, and PJRT boundary: backend ABI and lifecycle details.
• Feasibility and research: evidence and risks.
• Implementation roadmap: executable delivery slices.

One-Sentence Contract

SGraph turns a statically constructed local Graph into a transformed StableHLO module, compiles it through PJRT/XLA for an explicit replica mesh, and executes one local-buffer argument list per addressable replica.