[tunix] Add model-agnostic JAX diffusion training contracts

[ethannnnnn]

Motivation

Diffusion training integrations need a stable batch and scoring boundary that does not depend on MaxText, a specific model, or one training algorithm.

Scope

• Add a canonical batch-major DiffusionTokenBatch with target-aligned targets and explicit loss weights.
• Define typed raw-batch adapter and logits-function protocols.
• Validate shapes, dtypes, eager values, and scorer outputs while preserving JAX pytree behavior.

Design

Every model-input leaf has the same leading batch axis. Model-specific static configuration lives in the logits-function closure rather than the batch. Targets refer to physical target positions, so same-position and shifted models resolve alignment before constructing the contract.

Construction and scoring boundaries validate array structure and active values. Validation remains trace safe under JIT and sharding; additional finite/range checks run when arrays are eagerly addressable.

Compatibility

These are new model-agnostic JAX modules. Existing SFT, distillation, RL, and autoregressive paths do not construct them and remain unchanged.

Extensibility

MaxText and other integrations can provide different corruption, tokenizer, model, and alignment policies while sharing the same Tunix loss and learner interfaces.

Tests

• 10 focused diffusion contract tests at this stack point.
• Coverage includes invalid shapes/dtypes/values, target-aligned logits, pytrees, JIT, and sharding-safe validation.
• Pyink, isort, Pylint, Pyrefly, and Python compilation passed.

Known limitations

The contract intentionally does not create corruption, execute a model rollout, choose stop tokens, or prove that a prepared batch is fresh.

Stack

First PR in this repository stack; no preceding PR dependency.

Tunix block-diffusion design document

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.