记图与 jit

def f(x):
  return jnp.sum(jnp.sin(x) * x**2)

print(jax.make_jaxpr(f)(x1000))

fn f(x: &TracedTensor) -> R<TracedTensor> {
  Ok((&x.sin()? * &(x * x)?)?.reduce_sum(None)?)
}

let x = TracedTensor::input_concrete_shape(DType::F64, &[1000])?;
let y = f(&x)?;
let program = GraphCompiler::new().compile_with_input_specs(&y, &[(&x, DType::F64, &[1000][..])])?;
println!("inputs={} ops={}: {}", program.input_count(), op_names(&program).len(), op_names(&program).join(" "));

# JAX
{ lambda ; a:f64[1000]. let
    b:f64[1000] = sin a
    c:f64[1000] = integer_pow[y=2] a
    d:f64[1000] = mul b c
    e:f64[] = reduce_sum[axes=(0,) out_sharding=None] d
  in (e,) }
# tenferro
inputs=1 ops=4: Sin Mul Mul ReduceSum

jf = jax.jit(f)
v = jax.block_until_ready(jf(x1000))          # 第一次：追踪 + 编译 + 执行
# 之后 200 次取最小值；再用一个新函数对象单独量 lower() 和 compile()

let p = prepare(&rt, &eng, &program, &[&x1000])?;   // rt.prepare(&program, &reads, &[eng.target()])
let v = run1(&rt, &mut eng, &p, &[&x1000])?;        // rt.run(&mut [&mut eng], &p, &reads)

# JAX
value=1235.325872458312688 |.-exact|=0.0e+00
jit first call 19012 us, steady 7.5 us, traces so far 1
AOT on a copy: lower 1112 us, compile 16928 us, steady 9.3 us
# tenferro（进程里第一个程序）
build 1212.1 us, compile 534.0 us, prepare 1673.3 us, first run 4415.8 us, steady 6.5 us
prepare again (cache) 52.9 us; ...

v2 = jf(x2000)              # 同一个 jit 函数，长度 2000
jf(x1000)                   # 再回到 1000
compiled(x2000)             # AOT 编译出来的可执行文件

run1(&rt, &mut eng, &p, &[&x2000])               // 为 [1000] 准备的 Prepared
prepare(&rt, &eng, &program, &[&x2000])          // 为 [1000] 编译的程序
// 重建：新的占位符 [2000]、建图、编译、准备

# JAX
first call at 2000: 19597 us, value=2470.067398026252249 |.-exact|=0.0e+00, traces 2, cache size 2
back to 1000: 46 us, traces 2, cache size 2
AOT compiled(x2000): TypeError: Argument types differ from the types for which this computation was compiled. Perhaps you are calling the compiled executable with a different enable_x64 mode than when it was AOT compiled? The mismatches are:
# tenferro
run: Err: binding shape mismatch for concrete-shape placeholder: expected [1000], got [2000]
prepare: Err: binding shape mismatch for concrete-shape placeholder: expected [1000], got [2000]
rebuild for 2000: build+compile+prepare 85.2 us, value=2470.067398026252704

let xs_leaf = TracedTensor::from_tensor_symbolic_shape(Tensor::from_vec_col_major(vec![1000], xs(1000))?)?;
let ps = GraphCompiler::new().compile(&f(&xs_leaf)?)?;
let pp = prepare(&rt, &eng, &ps, &[&x1000])?;
// 同一个 pp 跑 1000 和 2000；再把求和换成 dot_general(sin(x), x*x) 试一遍

elementwise, prepared for 1000, run 1000: 1235.325872458312915
elementwise, prepared for 1000, run 2000: 2470.067398026252704
dot_general version ops: Sin Mul DotGeneral
dot_general, run 1000: 1235.325872458312915
dot_general, prepared for 1000, run 2000: Err: Runtime::run (Execution): runtime state failure: provider contract violation: prepared specialization projection mismatch
dot_general, prepare again for 2000 (32.9 us), run 2000: 2470.067398026252704

w = jnp.full(1000, 2.0)
g = lambda x: jnp.sum(jnp.sin(x) * w)
print(jax.make_jaxpr(g)(x1000))

{ lambda a:f64[1000]; b:f64[1000]. let
    c:f64[1000] = sin b
    d:f64[1000] = convert_element_type[new_dtype=float64 weak_type=False] a
    e:f64[1000] = mul c d
    f:f64[] = reduce_sum[axes=(0,) out_sharding=None] e
  in (f,) }

let w = TracedTensor::from_vec_col_major(vec![1000], vec![2.0_f64; 1000])?;
let xc = TracedTensor::input_concrete_shape(DType::F64, &[1000])?;
let yc = (&xc.sin()? * &w)?.reduce_sum(None)?;
let pc_prog = GraphCompiler::new().compile_with_input_specs(&yc, &[(&xc, DType::F64, &[1000][..])])?;

inputs=2 (x index Some(0), w index Some(1))
run with [x]: 1415.639514615922508

iteration 0: first call 17218 us, value 1415.639514615922508
iteration 1: first call 17242 us, value 2123.459271923883534
iteration 2: first call 18732 us, value 2831.279029231845016
w as an argument, iteration 0: 17617 us
w as an argument, iteration 1: 32 us
w as an argument, iteration 2: 24 us

cache before ... hits: 4, misses: 12, ...
cache after  ... hits: 6, misses: 12, ...

jax.jit(lambda x: x if jnp.sum(x) > 0 else -x)(x1000)

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].

println!("x placeholder attached_value: {}", x.attached_value().is_some());
println!("w data leaf attached_value: {}", w.attached_value().is_some());
println!("y = f(x) attached_value: {}", y.attached_value().is_some());
println!("w * w (derived from data only) attached_value: {}", (&w * &w)?.attached_value().is_some());

x placeholder attached_value: false
w data leaf attached_value: true
y = f(x) attached_value: false
w * w (derived from data only) attached_value: false

run with [x, w=3]: 2123.459271923883989
rebuilt with w=3: build+compile+prepare 47.3 us, value 2123.459271923883989