后端与集成

PyTorch 应用

PyTorch 是核心依赖。原生模型、挂载程序和文档限定的预训练流程可以进入同一个 ARTITrainingSession。按真实应用选择 CPU 或 CUDA Torch 构建，不能从 ARTI 包版本推断硬件就绪。

原生 JAX 程序

arti.jax 将共享 Formula 与有界/静态 ProgramGraph 声明绑定到原生数组。静态程序、参数 PyTree 与前向状态 PyTree 保持独立。编译调用运行 JAX 运算，不执行 Torch 或 host callback。

python -m pip install "arti-fit[jax,flax]==3.3.0"

• 在绑定的文档契约内使用 jit、grad 与 vmap。
• 固定长度 scan 携带明确返回的状态。
• 绑定按物理 shape 和图结构特化；自定义 Torch callable 或 VJP 需要明确的原生实现。

可选 Flax 集成

原生 Flax 集成遵循 init/apply，将参数与可变 collection 分开。声明的可变状态通过下一状态返回，不隐式成为优化器参数。原生挂载与参数 mask 让应用继续掌握宿主和适配器归属。

保留各适配器的范围

Accelerate、Transformers/TRL、Ray 与报告适配器接入各自文档限定的原生流程。上游损失、累积、mask、精度和执行语义仍然有效。存在适配器不等于所有模型或训练拓扑都经过验收。

阅读原生 lowering 限制

文档限定的 JAX 范围不包括前驱写回 Formula effect、任意 federated Query 程序的自动 lowering、开放长度反向模式或自动独立优化器区域推导。本地训练恢复不意味着原生 JAX 或分布式/分片 checkpoint 恢复。