elu: input_scale wrongly applied to the positive branch in forward and backward

[truong-v]

Describe the bug

elu returns wrong values whenever input_scale != 1, in both forward and backward. The positive branch of each kernel multiplies by input_scale, but ATen applies input_scale only to the negative branch. Errors are large (0.5 on the backward repro below), not a precision issue.

ATen (ActivationEluKernel.cu), with poscoef = scale, negcoef = alpha * scale, negiptcoef = input_scale:

forward:   x > 0 ? x * poscoef : expm1(x * negiptcoef) * negcoef
backward:  x > 0 ? grad * poscoef : grad * negiptcoef * negcoef * exp(x * negiptcoef)

Neither positive branch contains input_scale.

Reproduction

import torch
import flag_gems

torch.manual_seed(0)
x = torch.randn(8, device=flag_gems.device)
go = torch.ones_like(x)
alpha, scale, input_scale = 1.0, 1.0, 0.5

ref = torch.ops.aten.elu(x, alpha, scale, input_scale)
with flag_gems.use_gems():
    out = torch.ops.aten.elu(x, alpha, scale, input_scale)
print((out - ref).abs().max())  # 0.0726 — should be ~1e-7

ref = torch.ops.aten.elu_backward(go, alpha, scale, input_scale, False, x)
with flag_gems.use_gems():
    out = torch.ops.aten.elu_backward(go, alpha, scale, input_scale, False, x)
print((out - ref).abs().max())  # 0.5 — should be ~1e-7

The is_result=True backward path is wrong in the same way.

Root cause

Three positive branches in src/flag_gems/ops/elu.py carry a spurious input_scale factor:

line	current	should be
L32 elu_forward_kernel	scale * input_scale * x	scale * x
L45 elu_backward_kernel_with_self	grad_output * scale * input_scale	grad_output * scale
L58 elu_backward_kernel_with_result	grad_output * scale * input_scale	grad_output * scale

All three negative branches are already correct.

The forward defect dates to #415, the backward ones to #927.

Why the tests miss it

tests/test_elu.py hardcodes input_scale = 1.0, where the extra factor is a no-op. test_elu and test_elu_ also call torch.nn.functional.elu, which does not expose input_scale at all, so no existing case can reach the bug.

Fix

Drop input_scale from the three positive branches, and parametrize the tests over input_scale (calling torch.ops.aten.elu so the argument is reachable). Verified locally: 49 of 126 cases fail before the fix, all 126 pass after, max error vs ATen ~1e-7 in fp32. PR to follow.

Environment

• FlagGems master (bcf3f5b7)
• torch 2.6.0+cu124, triton 3.2.0, NVIDIA H100, CUDA 12.4

[saisharan0103]

The repro details point to input_scale leaking into the x > 0 path: with alpha=scale=1 and input_scale=0.5, forward drifts by 0.0726 and backward by 0.5 versus ATen. src/flag_gems/ops/elu.py is the part I would inspect first, especially the three positive branches called out in the report. For a PR, I would start by removing that factor from forward and both backward modes, then add a non-default input_scale regression test.