torch.tx.InferExpressionShape
export type InferExpressionShape<
Strings extends readonly string[],
Tensors extends readonly TensorOrNumber[],
> =
// Single tensor expression: -${x}, abs(${x}), etc.
Tensors extends readonly [infer T1]
? T1 extends TensorOrNumber
? ShapeOf<T1>
: DynamicShape
: // Binary operation: ${a} op ${b}
Tensors extends readonly [infer T1, infer T2]
? Strings extends readonly [string, infer Op extends string, string]
? T1 extends TensorOrNumber
? T2 extends TensorOrNumber
? BinaryOpShape<ShapeOf<T1>, ShapeOf<T2>, Op>
: DynamicShape
: DynamicShape
: DynamicShape
: // Complex expression (3+ tensors) → dynamic, checked at runtime
DynamicShape;Stringsextends readonly string[]Tensorsextends readonly TensorOrNumber[]Infer result shape from template strings and interpolated tensors.
Template structure analysis:
- 1 tensor:
${x}or-${x}→ preserves shape - 2 tensors:
${a} op ${b}→ binary operation shape - 3+ tensors: complex expression → DynamicShape (runtime checked)
Examples
// Template ["", " + ", ""] with tensors [Tensor<[2,3]>, Tensor<[2,3]>]
// → BroadcastShape<[2,3], [2,3]> = [2, 3]
// Template ["", " @ ", ""] with tensors [Tensor<[2,3]>, Tensor<[3,4]>]
// → MatmulShape<[2,3], [3,4]> = [2, 4]
// Template ["", " @ ", ""] with tensors [Tensor<[2,3]>, Tensor<[5,4]>]
// → matmul_error_inner_dimensions_do_not_match<3, 5, ...>