Rust in Everything
基于 tenferro-runtime 0.8.0

missing shape hint for traced tensor

主线位置:traced 主线的建图一步,在符号形状的输入上调激活、softmax、归一化。这些运算为什么是组合运算、怎么分解见 神经网络运算;四种叶子和符号形状见 TracedTensor。

flowchart TB
  x["符号形状输入"] --> prim["exp tanh 等"]
  prim --> ok1["建图通过"]
  x --> comp["sigmoid 等组合运算"]
  comp --> shape["ops.shape 要具体形状"]
  shape --> err["missing shape hint"]
  c["具体形状输入"] --> comp2["sigmoid"]
  comp2 --> ok2["建图通过"]

实验 crate exp-tenferro8/pit2-exp,bin p28_symbolic_composite,CPU(CpuEngine::serial()),debug,M2 Max、macOS 27、rustc 1.99.0。

现象

两种符号形状的输入:不带数据的占位符 input_symbolic_shape(DType::F64, 1),和带 3 个元素的符号叶子 from_tensor_symbolic_shape。在上面各调一遍:

for (op, r) in [
  ("exp", x.exp()),
  ("sigmoid", x.sigmoid()),
  ("silu", x.silu()),
  ("softplus", x.softplus()),
  ("gelu", x.gelu()),
  ("gelu_tanh", x.gelu_tanh()),
  ("softmax(0)", x.softmax(0)),
  ("layer_norm", x.layer_norm(0, None, None, 1e-5)),
] { /* 打印 Ok 或 Err */ }
placeholder exp: Ok
placeholder sigmoid: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 0)
placeholder silu: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 0)
placeholder softplus: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 4)
placeholder gelu: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 0)
placeholder gelu_tanh: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 0)
placeholder softmax(0): Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 0)
placeholder layer_norm: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 0)
symbolic leaf exp: Ok
symbolic leaf sigmoid: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 1)
...
symbolic leaf layer_norm: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 1)
  • exp 通过,七个组合运算全部失败,带不带数据都一样。
  • 消息里没有 sigmoid 这个名字,只有 TracedTensor::concrete_shape 和一个内部编号。softplus 报的是 tensor 4,一个中间值,用户代码里根本没有它。
  • 在同一个符号叶子上,tanh 能用,x * 0.5(乘一个秩 0 的常数)也报同一条错:
symbolic * rank-0 constant: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 1)
symbolic tanh: Ok -> [-0.7615941559557649, 0.0, 0.7615941559557649]

出处

组合运算定义一次、在 traced 和 eager 两个面上共用,接口是 CompositeOps(神经网络运算 的「定义一次:CompositeOps」)。它的 shape 方法的文档写明「Returns a ValidationError::InvalidArgument when the shape is symbolic」。sigmoid 第一步就要形状,用来把常数 1 广播成输入的形状:

    let shape = ops.shape(x)?;
    let (positive, neg_abs) = positive_and_neg_abs(ops, x)?;
    let e = ops.unary(CompositeUnary::Exp, &neg_abs)?;
    drop(neg_abs);
    let one = splat(ops, dtype, 1.0, &shape)?;

traced 一侧的 shape 落到 concrete_shape,形状提示里有符号维度就报「missing shape hint for traced tensor N」。softmax、归一化要拿形状算规约后的形状和广播,同理。softplus 先做了几步不需要形状的运算才碰到 shape,所以报的是中间值的编号。乘一个秩 0 常数失败是另一处:秩不同时 broadcast_binary 也要具体形状(add: incompatible shapes 的「出处」)。exp、tanh、erf 是 primitive,形状从输入原样传下去,不需要知道长度。

这不是 bug,是设计取舍:组合运算换来的是不用写 AD 规则、traced 程序里只有核心算子;代价就是建图时要具体形状。

修法

用具体形状的输入。要「编译一次、喂数据」,用 input_concrete_shape 建占位符,编译时给出同样的形状:

let c = TracedTensor::input_concrete_shape(DType::F64, &[3])?;
let y = c.sigmoid()?;
let prog = GraphCompiler::new().compile_with_input_specs(&y, &[(&c, DType::F64, &[3])])?;
let x = Tensor::from_vec_col_major(vec![3], vec![-1.0_f64, 0.0, 1.0])?;
println!("concrete placeholder sigmoid: {:?}", cpu.run_program(&prog, &[&x])?.as_slice::<f64>()?);
concrete placeholder sigmoid: [0.2689414213699951, 0.5, 0.7310585786300049]
  • 长度会变时,每种长度建一次图、编译一次;prepare 按结构和特化缓存,同一种长度只准备一次(执行与缓存)。
  • eager 下没有这个问题:eager 的张量总有具体形状,组合运算在不跟踪时还会走融合内核(神经网络运算)。
  • 只用 primitive 改写(例如把 sigmoid 写成 tanh 的式子)也绕不过去:式子里的常数 0.5、1 一样要广播,而秩 0 常数和符号张量相乘本身就需要具体形状。

怎么确认没中招

  • 搜 input_symbolic_shape 和 from_tensor_symbolic_shape,看它们之后有没有 sigmoid、silu、softplus、gelu、gelu_tanh、softmax、log_softmax、masked_softmax、layer_norm、rms_norm、take_along_axis、reduce_mean,以及和常数的二元运算。
  • 看到 TracedTensor::concrete_shape (GraphBuild),就往前找第一个符号形状的叶子。
  • XLA 后端的实验用的也是具体形状的输入,原因相同(XLA)。

小结

  • 组合运算(激活、softmax、归一化)在建图时要具体形状,符号形状的输入报 missing shape hint for traced tensor N,消息里没有运算名,N 可能是中间值。
  • exp、tanh 这类 primitive 不受影响;和秩 0 常数的二元运算同样需要具体形状。
  • 用 input_concrete_shape 或具体叶子,长度变了就重新建图编译;eager 不受影响。

💬 关联讨论

发起讨论
加载中...