Rust in Everything
基于 tenferro-runtime 0.8.0

dtype

主线位置：张量主线的第一步「建数组」时 dtype 从哪来，第二步「计算」里两种 dtype 相遇时结果是什么。
版本：JAX 0.12.0.1（CPU，x64 开和关两种）、NumPy 2.5.4、PyTorch 2.15.0（CPU，单线程）、tenferro-rs 0.8.0（CPU，native）
flowchart TB
  jx["JAX 数组"] --> flag{"x64 开了吗"}
  flag -->|否| f32["float32 int32"]
  flag -->|是| f64["float64 int64"]
  rs["Rust 字面量"] --> lit{"元素类型"}
  lit -->|"1.0"| F64["F64"]
  lit -->|"1"| I32["I32"]
  lit -->|"1.0_f32"| F32["F32"]
  F64 --> mix["混合 dtype"]
  F32 --> mix
  mix -->|"traced eager"| up["按提升格提升"]
  mix -->|"直接运算"| no["DTypeMismatch"]

实验代码：一个 JAX / NumPy / PyTorch 脚本（JAX 部分在 JAX_ENABLE_X64=0 和 =1 的子进程里各跑一遍，PyTorch torch.set_num_threads(1)），一个按两种 x64 设置各跑一遍的 JAX TRG 脚本（χ = 16、20 步，内置 SVD，只求值），一个 tenferro 程序 dtype（release，CpuEngine::serial()），以及试验阶段的 tenferro 分段 TRG 程序（重跑，单线程）。平台 M2 Max、macOS 27、rustc 1.99.0。

默认 dtype：一个开关，还是一个字面量

JAX 的默认 dtype 由全局开关 jax_enable_x64 决定，默认关：

print(jnp.array([1.0, 2.0]).dtype, jnp.array([1, 2]).dtype)
a = jnp.asarray(np.array([0.1, 0.2]))                 # NumPy float64 in
b = jnp.array([0.1, 0.2], dtype=jnp.float64)
s = float(jnp.sum(1.0 / jnp.asarray(x) ** 2))         # x = 1..1e5, float64

== JAX_ENABLE_X64=0
  jax_enable_x64 = False
  jnp.array([1.0, 2.0]).dtype = float32   jnp.array([1, 2]).dtype = int32
  jnp.asarray(np f64).dtype = float32  warnings: 0
  jnp.array(.., dtype=float64).dtype = float32  warnings: ['Explicitly requested dtype float64 requested in array is not available, and will be trunca']
  sum 1/k^2, k<=1e5 = 1.6449239253997803  err vs f64 = -1.4149844629329777e-07
== JAX_ENABLE_X64=1
  jax_enable_x64 = True
  jnp.array([1.0, 2.0]).dtype = float64   jnp.array([1, 2]).dtype = int64
  jnp.asarray(np f64).dtype = float64  warnings: 0
  jnp.array(.., dtype=float64).dtype = float64  warnings: []
  sum 1/k^2, k<=1e5 = 1.6449240668982266  err vs f64 = 0.0

开关关着时，NumPy 的 float64 数组经 jnp.asarray 进来就变成 float32，一条警告都没有；只有显式写了 dtype=float64 才警告。对 TRG 来说，这一步把结果带偏多少：

x64 = False
beta=0.30 dtype=float32 lnZ/N=0.790559232234955 err vs Onsager=+1.613e-07
beta=0.44 dtype=float32 lnZ/N=0.928723871707916 err vs Onsager=-4.669e-06
beta=0.60 dtype=float32 lnZ/N=1.210132718086243 err vs Onsager=+3.298e-07
...
x64 = True
beta=0.30 dtype=float64 lnZ/N=0.790559063800966 err vs Onsager=-7.150e-09
beta=0.44 dtype=float64 lnZ/N=0.928723962577071 err vs Onsager=-4.578e-06
beta=0.60 dtype=float64 lnZ/N=1.210133028164417 err vs Onsager=+6.399e-07

（Onsager 精确值由 SciPy 的 quad 积分得到，两次运行各打印一条 IntegrationWarning。）β = 0.30 时 χ = 16 的截断误差只有 7.2e-9，float32 把误差放大到 1.6e-7，大了二十多倍；β = 0.44、0.60 时截断误差本身更大，float32 让 lnZ/N 分别移动了 9.1e-8 和 3.1e-7。同样 20 步的 tenferro 分段实现（F64）给出 0.790559063800966、0.928723962577072、1.210133028164417，和 x64 的 JAX 相同或只差末位（完整对照见动手：Ising TRG）。

tenferro 没有这样的开关。张量的 dtype 就是构造它的 Rust Vec 的元素类型：

Tensor::from_vec_col_major(vec![2], vec![1.0, 2.0])?.dtype()
Tensor::from_vec_col_major(vec![2], vec![1, 2])?.dtype()
Tensor::from_vec_col_major(vec![2], vec![1_i64, 2])?.dtype()
Tensor::from_vec_col_major(vec![1], vec![1.0_f32])?.dtype()

vec![1.0, 2.0] -> F64
vec![1, 2]     -> I32
vec![1_i64, 2] -> I64
vec![1.0_f32]  -> F32
  • 浮点字面量没有后缀时 Rust 推成 f64，所以 tenferro 的「默认」相当于 JAX 开了 x64。不存在「忘了开开关」这回事。
  • 整数字面量没有后缀时 Rust 推成 i32，这一点和 JAX x64 的 int64 不同。下标张量用 I32 没问题（take_along_axis 接受 I32 和 I64），但要存大于 2³¹ 的计数时要写 _i64。
  • 带类型参数的 TypedTensor<f64> 把 dtype 放进 Rust 类型里，混用在编译期就是 E0308，见数据模型。

两个数组相遇：提升格

同形状的两个零数组相加，看结果的 dtype（JAX 开 x64）：

for a, b in pairs:
  print((jnp.zeros(2, dtype=a) + jnp.zeros(2, dtype=b)).dtype, ...)   # also numpy, torch

for (a, b) in pairs {
  let ta = TracedTensor::from_tensor_concrete_shape(zeros(a)?)?;
  let tb = TracedTensor::from_tensor_concrete_shape(zeros(b)?)?;
  let traced = (&ta + &tb)?;                                    // graph build, then eval
  let eager = ert.with_eager_session(&mut eng, |s| s.add(&ea, &eb));
  let direct = zeros(a)?.add(&zeros(b)?, &mut eng);             // TensorSessionOpsExt
}

# JAX / NumPy / PyTorch
     int32 + float32   jax float32    numpy float64    torch float32
     int64 + float32   jax float32    numpy float64    torch float32
     int32 + float64   jax float64    numpy float64    torch float64
   float32 + float64   jax float64    numpy float64    torch float64
   float64 + complex64 jax complex128 numpy complex128 torch complex128
      bool + int32     jax int32      numpy int32      torch int32
     int32 + int64     jax int64      numpy int64      torch int64
# tenferro
I32 + F32: traced F64 (run F64); eager F64; direct Err(add: dtype mismatch: expected I32, actual F32)
I64 + F32: traced F64 (run F64); eager F64; direct Err(add: dtype mismatch: expected I64, actual F32)
I32 + F64: traced F64 (run F64); eager F64; direct Err(add: dtype mismatch: expected I32, actual F64)
F32 + F64: traced F64 (run F64); eager F64; direct Err(add: dtype mismatch: expected F32, actual F64)
F64 + C32: traced C64 (run C64); eager C64; direct Err(add: dtype mismatch: expected F64, actual C32)
Bool + I32: traced I32 (run I32); eager I32; direct Err(add: dtype mismatch: expected Bool, actual I32)
I32 + I64: traced I64 (run I64); eager I64; direct Err(add: dtype mismatch: expected I32, actual I64)
I32.cast(F32) + F32: traced F32

（tenferro 的 C32 / C64 是 32 / 64 位浮点的复数，对应 complex64 / complex128。）

  • 七对里六对四边一致。唯一的分歧是整数加 float32：JAX 和 PyTorch 给 float32，tenferro 和 NumPy 给 64 位浮点。
  • tenferro 的理由写在 shape_infer::promote_dtype 的文档里：「I32/I64 + F32/F64 → F64 (F32 does not have enough mantissa bits for arbitrary integers)」。实际的格在 tenferro-tensor-core 的 promote_specs：整数遇上浮点，目标是「最宽的浮点」。
  • 后果：float32 的数据和一个整数张量（计数、下标算出来的权重）相乘，整条链变成 F64，内存和时间都翻倍，而且没有提示。要留在 F32，先把整数 cast(DType::F32)，最后一行就是这样。
  • traced 在建图时、eager 在执行前按这张格插入转换；直接运算（Tensor 在 engine 上的方法）不提升，直接报 DTypeMismatch。这一条的来龙去脉见常见坑 dot_general: dtype mismatch。

弱类型：Python 标量和 scale_real

JAX 里 Python 的 float 是「弱类型」：它让位给数组的 dtype，float32 数组乘 0.1 还是 float32；换成 NumPy 的 np.float64(0.1) 就是强类型，结果提升成 float64：

x32 = jnp.array([1.0, 3.0], dtype=jnp.float32)
print(x32 * 0.1, x32 * np.float64(0.1))

  f32 * 0.1             float32 [0.10000000149011612, 0.30000001192092896]
  f32 * np.float64(0.1) float64 [0.1, 0.30000000000000004]
...
  bits of (f32 [1, 3] * 0.1): ['0x3dcccccd', '0x3e99999a']

tenferro 的运算只在张量之间进行，没有「Python 标量」这一类。对应弱类型的是 scale_real(f64)：系数按张量的 dtype 落成常量，dtype 不变。对应强类型的是一个真正的 F64 张量，它按提升格把结果变成 F64：

let x32 = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f32, 3.0])?;
let s = x32.scale_real(0.1)?;
let c = TracedTensor::from_vec_col_major(vec![], vec![0.1_f64])?;
let p = (&x32 * &c.broadcast_in_dim(&[2], &[])?)?;

F32 .scale_real(0.1): F32 [0.1, 0.3] bits ["0x3dcccccd", "0x3e99999a"]
F32 * F64 tensor 0.1: F64 [0.1, 0.30000000000000004]
TypedTensor<f32> scale_real(0.1) direct: [0.1, 0.3]

两边的 float32 结果逐位相同。TracedTensor::scale_real 对 F32 写的是 StdTensorOp::constant(factor as f32)，先把 0.1 舍入成 f32 再乘，和 JAX 对弱类型标量的处理一样。同一个函数对整数 dtype 把系数四舍五入成整数，那是另一回事，见整数运算。

astype 拆成两个：convert 和 cast

JAX 的 astype 什么方向都接受，变窄、浮点转整数都静默完成。tenferro 拆成两个方法：convert 只接受提升格允许的方向（变宽），cast 才做有损转换：

vals = [-2.7, 2.7, 2.5, -0.5, 1e20, -1e20, float("nan"), float("inf")]
print(jnp.array(vals).astype(jnp.int32), jnp.array([0.1, 1e300]).astype(jnp.float32))

let v = TracedTensor::from_vec_col_major(vec![vals.len()], vals.clone())?;
v.convert(DType::F32);     // Err at graph build
eval(&v.cast(DType::I32)?) // Err at run

# JAX（NumPy 的 int32 / int64、PyTorch 的 int64 在这台机器上给出同样的数）
  input               [-2.7, 2.7, 2.5, -0.5, 1e+20, -1e+20, nan, inf]
  jnp astype(int32)   [-2, 2, 2, 0, 2147483647, -2147483648, 0, 2147483647]
  jnp astype(int64)   [-2, 2, 2, 0, 9223372036854775807, -9223372036854775808, 0, 9223372036854775807]
...
  jnp f64 1e300 astype(float32): [inf]
# tenferro
traced cast I32: Err(cast: invalid argument value: real value 100000000000000000000 is out of i32 range)
...
traced convert F32: Err(TracedTensor::convert: unsupported dtype conversion from F64 to F32: checked convert only accepts conversions allowed by dtype promotion; use explicit cast for lossy dtype projection)
direct convert F32: Err(convert: unsupported dtype conversion from F64 to F32: checked convert only accepts conversions allowed by dtype promotion; use explicit cast for lossy dtype projection)
traced convert I64: Err(TracedTensor::convert: unsupported dtype conversion from F64 to I64: checked convert only accepts conversions allowed by dtype promotion; use explicit cast for lossy dtype projection)
traced cast I32 [-2.7, 2.7, 2.5, -0.5]: [-2, 2, 2, 0]
traced cast I32 [NaN]: Err(cast: invalid argument value: real value must be finite when casting to i32, got NaN)
traced cast I32 [inf]: Err(cast: invalid argument value: real value must be finite when casting to i32, got inf)
traced cast I32 [2147483648.0]: Err(cast: invalid argument value: real value 2147483648 is out of i32 range)
compile + prepare of cast(1e20 -> I64): ok
run: Err TensorRuntime(Validation { op: "cast", source: InvalidArgument { argument: "value", message: "real value 100000000000000000000 is out of i64 range" } })
traced cast F32 [0.1, 1e300]: [0.1, inf] bits of 0.1 0x3dcccccd
  • 范围内的值两边一样，都向零截断：−2.7 → −2，2.5 → 2，−0.5 → 0。
  • 范围外的值：JAX 饱和到 int32 的上下界，NaN 变成 0。tenferro 的 cast 不饱和，整个运算返回错误。检查在 CPU 后端的 validate_real_cast_to_i32：不是有限值、或超出 i32::MIN..=i32::MAX，就报 InvalidArgument。它要看数据，所以发生在执行时：编译和准备都成功，run 返回 Err。
  • cast 到 F32 不检查范围，1e300 变成 inf，和 JAX 一样；0.1 舍入出的位模式也相同（0x3dcccccd）。
  • convert 拒绝的是方向，不是数值：F64 → F32、F64 → I64 都在建图时报错，direct 运算的 convert 也一样（报 convert: 前缀）。这条规则来自规格 tensor-semantics.md 的「DType conversion」一节：convert 是「checked」，cast 是「explicit」。
📝NOTE
没有蠢问题

问： 我在 JAX 里一直开着 x64，移植到 tenferro 之后，哪些地方的 dtype 会和原来不一样？

答： 按上面的实测，三处。一是整数字面量：vec![1, 2] 是 I32，不是 int64。二是整数和 float32 混算：JAX 是 float32，tenferro 是 F64。三是浮点转整数时的越界值：JAX 饱和，tenferro 报错。其余的组合（F32 + F64、F64 + C32、Bool + I32、I32 + I64）以及 float32 乘一个实数系数的结果，两边相同。

小结

  • JAX 的默认 dtype 看 jax_enable_x64，关着时 float64 输入经 jnp.asarray 静默变成 float32，TRG 在 β = 0.3 的误差因此从 7e-9 涨到 1.6e-7；tenferro 的 dtype 就是 Rust 字面量的类型，1.0 是 F64，1 是 I32。
  • 提升格和 JAX 只在一处不同：整数加 float32，tenferro（和 NumPy）给 F64，JAX 和 PyTorch 给 float32。traced 和 eager 自动提升，直接运算报 DTypeMismatch。
  • tenferro 没有弱类型；scale_real 按张量 dtype 落系数，对 float32 的效果和 JAX 的 Python 标量逐位相同，乘一个 F64 张量则提升成 F64。
  • astype 对应两个方法：convert 只许变宽，cast 做有损转换；浮点转整数遇到 NaN、无穷或越界值时 tenferro 在执行时报错，JAX 饱和。

💬 关联讨论

发起讨论
加载中...