JAX 一侧静默降成 float32，对不上

主线位置：张量主线的第一步「建数组」，不过出事的是 JAX 那一侧：拿 JAX 写对照用的参考实现时。

版本：JAX 0.12.0.1（CPU，x64 开和关两种）、tenferro-rs 0.8.0（CPU，native）

flowchart TB
  np["NumPy f64 数据"] --> off["x64 关"]
  np --> on["x64 开"]
  off --> f32["jnp 数组 f32"]
  on --> f64["jnp 数组 f64"]
  f32 --> d1["差 4e-7"]
  f64 --> d2["差 4e-16"]
  tf["tenferro F64"] --> d1
  tf --> d2

实验代码：tenferro 程序 p08_grad_file（traced ad.grad，CpuEngine::serial()，release）把梯度写成文本；JAX 脚本在 JAX_ENABLE_X64=0 和 =1 两个子进程里各算一次同一个梯度再和它比。函数是变换主线的 f(x) = Σ sin(x)·x²，n = 1000，x_i = −2 + 4i/999。平台 M2 Max、macOS 27、rustc 1.99.0。

现象

tenferro 的梯度是 F64：

grad dtype F64, n = 1000, grad[:3] = [1.9726023611141572, 1.9931199187584416, 2.0134492414958083]

JAX 参考实现照常写，数据来自 NumPy（float64）：

x_np = np.array([-2 + 4 * i / (N - 1) for i in range(N)])
x = jnp.asarray(x_np)
g = jax.jit(jax.grad(lambda x: jnp.sum(jnp.sin(x) * x**2)))(x)

忘了开 x64 时：

== JAX_ENABLE_X64=0
  jax_enable_x64=False  x.dtype=float32  grad.dtype=float32  warnings=0
  max|jax - tenferro| = 4.125e-07  max|jax - exact| = 4.125e-07  max|grad| = 3.153
  jnp.array(.., dtype=float64).dtype=float32  warnings=['Explicitly requested dtype float64 requested in array is not available, and will']
  allclose(rtol=1e-12, atol=1e-14): False

jnp.asarray 把 float64 的 NumPy 数组静默降成 float32，0 条警告；只有显式写 dtype=jnp.float64 的那一行警告了一次（第三行，警告文字被截短）。两边差 4.1e-7，比 f64 对照该有的 1e-15 量级大八个数量级；按 f64 的容差判断，对照失败。乍看像 tenferro 的梯度错了，其实差的全部来自 JAX 一侧（max|jax - exact| 同样是 4.1e-7）。

出处

JAX 的默认 dtype 由全局开关 jax_enable_x64 决定，默认关：浮点默认 float32，传进来的 float64 也降成 float32。只有显式要 float64 时才警告，从 NumPy 转过来时不警告。tenferro 没有这个开关，dtype 来自 Rust 的元素类型，vec![1.0_f64] 就是 F64。两边默认值的并排对照见 dtype；全卷的对照为什么都开 x64，见怎么读这一卷。

修法

在脚本最开头、建任何数组之前开 x64（或在环境变量里设 JAX_ENABLE_X64=1）：

import jax
jax.config.update("jax_enable_x64", True)

== JAX_ENABLE_X64=1
  jax_enable_x64=True  x.dtype=float64  grad.dtype=float64  warnings=0
  max|jax - tenferro| = 4.441e-16  max|jax - exact| = 4.441e-16  max|grad| = 3.153
  jnp.array(.., dtype=float64).dtype=float64  warnings=[]
  allclose(rtol=1e-12, atol=1e-14): True

差 4.4e-16（量级 3.2），是 f64 舍入的量级；剩下的差别来自 sin、cos 的实现和 jit 的 FMA 融合，怎么读这一卷有分析。

开关要在建数组之前打开。打开之前建好的数组留在 float32：

== config.update after the first array
  before update: float32  after update: float64  a + b: float64

a 是打开开关之前从同样的 float64 数据建的，已经丢了精度；之后和 float64 的 b 相加，结果的 dtype 是 float64，看上去一切正常，丢掉的位数却回不来。

怎么确认没中招

• 对照脚本里加断言：assert x.dtype == jnp.float64，对输入和输出各一次。
• 对照的容差按 f64 定（这一卷用 1e-14·max(1, |ref|) 量级），差到 1e-7、1e-8 量级时，先查 JAX 一侧的 dtype，再查 tenferro。
• 从 .npy、.npz 读数据时更要小心：文件里是 float64，进了 jnp 就变了。

小结

• JAX 不开 x64 时，NumPy 的 float64 数据进 jnp.asarray 静默变成 float32，没有警告。
• 和 tenferro 的 F64 结果一比，差 1e-7 量级，看起来像 tenferro 错了。
• 脚本开头 jax.config.update("jax_enable_x64", True)，并断言输入输出的 dtype。