diff --git a/src/base/elu.h b/src/base/elu.h index 0d942f97f..9d47fd49e 100644 --- a/src/base/elu.h +++ b/src/base/elu.h @@ -7,6 +7,21 @@ namespace infini::ops { class Elu : public Operator { public: + Elu(const Tensor input, const double alpha, Tensor out) + : input_shape_{input.shape()}, + input_strides_{input.strides()}, + input_type_{input.dtype()}, + out_shape_{out.shape()}, + out_strides_{out.strides()}, + out_type_{out.dtype()}, + alpha_{alpha}, + scale_{1.0}, + input_scale_{1.0}, + device_index_{out.device().index()} {} + + /// \deprecated Use `Elu(input, alpha, out)`. This constructor will be + /// removed in a future release. + [[deprecated("Use the `(input, alpha, out)` overload instead.")]] Elu(const Tensor input, const double alpha, const double scale, const double input_scale, Tensor out) : input_shape_{input.shape()}, @@ -20,6 +35,13 @@ class Elu : public Operator { input_scale_{input_scale}, device_index_{out.device().index()} {} + void operator()(const Tensor input, const double alpha, Tensor out) const { + return operator()(input, alpha, 1.0, 1.0, out); + } + + /// \deprecated Use `operator()(input, alpha, out)`. This overload will be + /// removed in a future release. + [[deprecated("Use `operator()(input, alpha, out)` instead.")]] virtual void operator()(const Tensor input, const double alpha, const double scale, const double input_scale, Tensor out) const = 0; diff --git a/tests/test_torch_ops.py b/tests/test_torch_ops.py index 99021467d..2bfd8c50d 100644 --- a/tests/test_torch_ops.py +++ b/tests/test_torch_ops.py @@ -96,8 +96,8 @@ ("softplus", "beta"): 1.0, ("softplus", "threshold"): 20.0, ("elu", "alpha"): 1.0, - ("elu", "scale"): 1.0, - ("elu", "input_scale"): 1.0, + ("elu", "scale"): 1.25, + ("elu", "input_scale"): 0.75, ("sub", "alpha"): 1.0, ("addcmul", "value"): 1.0, ("addcdiv", "value"): 1.0, @@ -530,8 +530,13 @@ def test_op(op_meta, shape, dtype, device, rtol, atol): clone_strided(x) if isinstance(x, torch.Tensor) else x for x in inputs ] + if aten_name == "elu": + ref_func = torch.ops.aten.elu.default + else: + ref_func = _torch_func(aten_name) + try: - ref = _torch_func(aten_name)(*ref_inputs) + ref = ref_func(*ref_inputs) except ( RuntimeError, TypeError, @@ -595,3 +600,11 @@ def test_op(op_meta, shape, dtype, device, rtol, atol): for actual, expected in zip(outs, ref_outs): _assert_close(actual, expected, rtol, atol) + + if aten_name == "elu": + input, alpha = inputs[:2] + alpha_out = torch.empty_like(ref_outs[0]) + _call_infini(op_name, input, alpha, alpha_out) + _assert_close( + alpha_out, torch.nn.functional.elu(input, alpha=alpha), rtol, atol + )