[tunix] Add model-agnostic JAX diffusion contracts

[ethannnnnn]

Motivation

Diffusion training integrations need a small, stable JAX boundary that does
not depend on MaxText, a specific model, or a training algorithm.

Scope

This PR adds only:

• a batch-major DiffusionTokenBatch with target-aligned token IDs and
  explicit loss weights;
• typed raw-batch adapter and logits-function protocols;
• shape, dtype, eager-value, and scorer-output validation; and
• focused contract tests and package documentation.

It does not add an SFT or distillation objective, change gradient reduction,
prepare a rollout, alter RL behavior, or import MaxText.

Design

Every model-input leaf has the same leading batch axis. Static model
configuration remains in the logits-function closure rather than becoming a
batch leaf. Targets refer to physical prediction positions, so model-specific
same-position or shifted-logit alignment happens before constructing the
canonical batch.

DiffusionLogitsFn receives an NNX model and prepared model inputs and returns
floating-point logits with shape [batch, length, vocabulary]. Construction
and scoring boundaries validate array structure, target alignment, and active
values. Static checks remain safe under JIT and sharding; concrete finite and
range checks run when arrays are eagerly addressable.

Compatibility

The package is new and opt-in. Existing SFT, distillation, RL, and
autoregressive paths do not construct these contracts and remain unchanged.

Tests

python -m pytest -q tests/diffusion/types_test.py
15 passed

Pyink, isort, scoped Pylint, Python compilation, and git diff --check also
pass for the changed files.

Follow-Ups

Denominator-aware weighted reduction will be proposed separately after this
contract lands. Diffusion SFT, prepared distillation, OPD, rollout transport,
policy scoring, and GRPO integration remain later independent reviews.

Design document:
https://docs.google.com/document/d/1Xe-98ScS2RSH29AdhTdc9tO4WHQG5wCbFkJA3gmIUZQ/edit