fix: teleported operators precision

[AntoineF4C5]

Improves precision of teleported op by improving the way the corresponding lookup table is constructed.

Context

The main idea of teleported operators is to first reduce the domain of the input.
For divided operators, we are doing that by dividing the input by a constant τ.

We then lookup the divided input into a preconstructed table LUT.

Previous flow

Previously, the LUT was constructed as follow: LUT[x] = intended_op(x).
The issue was that, when taking into account the previous division, this lead us roughly to f(x) = LUT[x/τ] = intended_op(x/τ), where intended_op(x/τ) != intended_op(x).

Proposed solution

We now "reciprocate" the teleportation, by computing the LUT the following way: LUT[x] = intended_op(x * τ), allowing us to compute f(x) = LUT[x/τ] = intended_op( (x/τ)*τ ), where (x/τ)*τ roughly equals τ, with an error e<τ.

--

Closes #101

[Copilot]

Pull request overview

This PR improves the numerical accuracy of neural-teleport “divided” activation operators by constructing activation LUTs as op(x * τ) and reciprocating the teleportation division during execution so the effective input is closer to the original value.

Changes:

• Update neural-teleport signed activation LUT materialization to evaluate activations at (index * τ).
• Thread tau into TanhTable/SigmoidTable/ErfTable construction on the proof side.
• Update tracer implementations of tanh/sigmoid/erf to multiply the teleported quotient back by τ before evaluating the activation, and un-ignore related precision tests.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
jolt-atlas-core/src/onnx_proof/ops/tanh.rs	Passes tau into tanh LUT construction for proofs/verification.
jolt-atlas-core/src/onnx_proof/ops/sigmoid.rs	Passes tau into sigmoid LUT construction for proofs/verification.
jolt-atlas-core/src/onnx_proof/ops/erf.rs	Passes tau into erf LUT construction for proofs/verification.
jolt-atlas-core/src/onnx_proof/neural_teleport/utils.rs	Changes LUT materialization semantics to use (index * tau); extends LUT structs to store tau.
atlas-onnx-tracer/src/tensor/mod.rs	Adds scalar Tensor * scalar multiplication support used by teleport reciprocation.
atlas-onnx-tracer/src/ops/tanh.rs	Reciprocates teleportation (non-fused path) by multiplying by tau; re-enables precision test.
atlas-onnx-tracer/src/ops/sigmoid.rs	Reciprocates teleportation by multiplying by tau; re-enables precision test.
atlas-onnx-tracer/src/ops/erf.rs	Reciprocates teleportation by multiplying by tau; re-enables precision test.

Comments suppressed due to low confidence (1)

atlas-onnx-tracer/src/ops/tanh.rs:22

• With feature = "fused-ops" enabled, this implementation still computes tanh on the divided value (x/τ) via the LUT lookup, while the non-fused path now reciprocates teleportation (tanh(floor(x/τ)*τ)). This creates behavior/accuracy divergence between feature flags and undermines the PR’s stated precision fix; consider applying the same * τ reciprocation before the LUT lookup (or regenerating the LUT to match the new semantics).

        let input = tensor::ops::nonlinearities::const_div(inputs[0], self.tau as f64);

        #[cfg(feature = "fused-ops")]
        {
            let scale_i = self.scale.0 as i64;
            let lut = generate_tanh_lut(scale_i);
            input
                .par_enum_map(|_, a_i| {
                    Ok::<_, TensorError>(tanh_lut_lookup(a_i, scale_i as i32, &lut))
                })
                .unwrap()

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.