[Feature]: Define graph-level RNG key semantics for future random ops

[shinaoka]

What are you trying to do?

Capture the design discussion needed before tenferro-rs adds random-number generation to eager or traced tensor APIs. Future RNG work should have explicit graph-level semantics rather than inheriting hidden host state or backend-specific behavior by accident.

The target workflow is a JAX-like traced graph where random values are reproducible, compilable, backend-portable, and usable under batching/parallel execution.

What makes it hard today?

tenferro-rs does not currently expose a supported random tensor API. Adding one would affect public API design, primitive vocabulary, traced graph lowering, CPU/GPU backend behavior, cache keys, documentation, and AD semantics.

The main design risk is treating randomness as an implicit side effect. A hidden global RNG state would be hard to make reproducible across compile boundaries, graph optimization, backend dispatch, multi-device execution, and replay/oracle tests.

Proposed behavior or API direction

Use JAX's classic functional PRNG model as the primary reference point:

• Represent RNG state as an explicit graph value, usually called a key.
• Random samplers are pure operations of (key, shape/dtype/distribution params) and do not mutate the key.
• Reusing the same key intentionally produces the same random output; new randomness requires explicit key derivation.
• Include explicit key operations such as seed/key creation, split, and fold-in.
• Lower samplers through a raw random-bits primitive plus deterministic transforms where practical.
• Do not promise NumPy-style sequential equivalence: sampling N values in one op need not equal splitting into N scalar draws.
• Prefer a software counter-based/splittable PRNG for CPU/GPU portability and reproducibility; avoid hardware RNG semantics as the first design.
• Treat stateful RNG, if ever added, as a later convenience wrapper over explicit key/counter state, not as the base graph primitive model.

Possible primitive vocabulary to design:

• RandomSeed or KeySeed: seed integer -> key
• RandomSplit: key -> key array / multiple keys
• RandomFoldIn: key + 32-bit data -> derived key
• RandomBits: key + bit width + shape -> unsigned integer tensor
• Public samplers such as uniform, normal, randint, etc. built above the primitive boundary

Relevant area

• tenferro-core-ops: primitive catalog / public operation vocabulary
• tenferro-internal-ops: graph primitive semantics and AD rule manifest
• tenferro-runtime: traced graph construction, metadata, shape inference, lowering, execution IR/cache keys
• tenferro-tensor: CPU backend reference implementation
• tenferro-gpu: eventual CUDA backend parity and deterministic behavior
• tenferro-ad: explicit non-differentiability of keys/random-bit generation, and behavior of sampled values in downstream differentiable graphs
• docs/design/, docs/spec/, and user-facing docs

Dependency, backend, AD, and cache impact

• This likely introduces a new public operation family, so it must start as an accepted feature/design issue before implementation.
• If a PRNG implementation dependency is needed, it should be justified and added through workspace dependencies if shared.
• CPU and GPU behavior should not silently diverge. If exact parity is required, document and test it; if not, define the supported reproducibility contract precisely.
• RNG keys and raw random-bit generation should not be differentiable with respect to the key. Sampled tensors can still participate in later differentiable computation as ordinary primal values.
• Static parameters such as distribution type, shape, dtype, bit width, and PRNG implementation must participate in graph/runtime cache identity. Sample values must not be embedded in cache keys.
• Avoid hidden materialization and repeated per-element index work in CPU/GPU kernels.

Non-goals for the first design pass

• Do not implement the full RNG API directly from this issue without maintainer acceptance.
• Do not add implicit global RNG state as the primary graph semantics.
• Do not promise NumPy/PyTorch-style stateful generator behavior for traced graphs.
• Do not rely on backend hardware RNG as the default implementation.
• Do not add stateful RNG to transformed graph code in the first pass; track that separately if needed.

Acceptance criteria

• A design/spec document records tenferro's graph-level RNG contract, including key ownership/reuse, split/fold-in semantics, reproducibility scope, backend parity expectations, and non-goals.
• The primitive catalog or equivalent design notes identify the first RNG primitive set and where public wrappers should live.
• AD behavior is specified: key operations and random-bit generation are non-differentiable; downstream sampled numeric tensors behave like ordinary primal inputs to later ops.
• Shape/dtype/static-parameter rules are specified for traced execution and cache identity.
• Tests are planned for at least deterministic repeat with the same key, different outputs from split keys, no key mutation, fold-in behavior, batching behavior, and documented lack of sequential equivalence.
• If GPU support is included in the implementation phase, CPU/GPU reproducibility or divergence expectations are tested and documented.
• Stateful RNG is either explicitly deferred or captured in a separate issue with its own effect/ref semantics.

Notes and related references

JAX references from the design discussion:

• JAX pseudorandom numbers: https://docs.jax.dev/en/latest/random-numbers.html
• jax.random module: https://docs.jax.dev/en/latest/jax.random.html
• JAXPR language: https://docs.jax.dev/en/latest/jaxpr.html
• JEP 263, JAX PRNG design: https://docs.jax.dev/en/latest/jep/263-prng.html
• JEP 9263, typed keys and pluggable RNGs: https://docs.jax.dev/en/latest/jep/9263-typed-keys.html
• JEP 28845, optional stateful randomness: https://docs.jax.dev/en/latest/jep/28845-stateful-rng.html

Repository references:

• docs/reference/stablehlo-primitives.md already notes StableHLO rng and rng_bit_generator as compatibility-oriented ops.
• docs/reference/jax-primitives.md lists JAX RNG-related primitive references.
• REPOSITORY_RULES.md requires new public operation families, dependencies, feature flags, backend behavior, or AD semantics changes to start from issue intake.

Provenance

This issue records maintainer/agent design discussion about how JAX handles random numbers in graphs. No prototype implementation code is attached, and no external code should be copied from the linked references without a separate license/provenance review.