JAX 函数式后端

ARTI 仍以 PyTorch 为第一后端。arti.jax 是一个刻意收窄的原生 JAX 函数式子集，不是 fitting、挂载、序列化、Recall 或完整模块栈的第二套实现。

运行状态与 smoke gate

import arti.jax as aj

print(aj.backend_status())          # available | broken | unavailable
report = aj.smoke_report()
print(report["smoke_status"])

Smoke report 检查 eager 执行、JIT、整个参数 PyTree 的梯度以及 VMAP 语义。通过只证明这里列出的子集。

函数式 Layer

import jax
import jax.numpy as jnp
import arti.jax as aj

params = aj.init_layer(
    jax.random.PRNGKey(0), input_dim=32, hidden_dim=64, coord_dim=8
)
x = jnp.ones((4, 16, 32))
coord = jnp.ones((4, 16, 8))
mask = jnp.ones((4, 16), dtype=bool)
out = aj.apply_layer(params, x, coord=coord, mask=mask)

params 是只包含数组的 PyTree。需要显式 jax.vmap 组合时，可使用单样本形式 apply_layer_single()。

共享张量契约

当前子集包含 masked_mean、masked_softmax、mask_coverage、ensure_visibility、attention_mask_to_visibility 与 apply_coord_frame_inverse。测试会将选定的 mask、visibility 和 operator-bank 框架逆变换行为与 PyTorch 函数后端对照。

边界

这里不支持 arti.fit()、模型挂载、Recall 与 Stateful Recall、Membrane、Half/Fold/Pulse、checkpoint、多设备执行和完整 ARTILayer parity。状态报告会区分 JAX 未安装与本地 jaxlib 运行时损坏。