diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 44536e675da..e3a636466c1 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -163,6 +163,8 @@ from executorch.exir.dialects._ops import ops as exir_ops from torch.fx.node import Node +_LEAKY_RELU_DEFAULT_NEGATIVE_SLOPE = 0.01 + def require_static_int(value: Any, param_name: str, op_name: str) -> None: """ @@ -2786,6 +2788,63 @@ def _relu_handler(P: MLXProgramBuilder, n: Node) -> Slot: return out +@REGISTRY.register(target=[torch.ops.aten.leaky_relu.default]) +def _leaky_relu_handler(P: MLXProgramBuilder, n: Node) -> Slot: + """Handle aten.leaky_relu.default - leaky rectified linear unit. + + leaky_relu(x) = x if x >= 0 + = slope * x otherwise + + Implemented as where(x >= 0, x, slope * x) so it stays correct for any + negative_slope (including values > 1), matching eager PyTorch. + """ + args = P.args(n) + require_args(args, 1, 2, "aten.leaky_relu") + require_kwargs(P.kwargs(n), set(), "aten.leaky_relu") + + x = args[0] + negative_slope = _LEAKY_RELU_DEFAULT_NEGATIVE_SLOPE + if len(args) > 1 and args[1] is not None: + negative_slope = float(args[1]) + + x_meta = n.args[0].meta.get("val") + if x_meta is None: + raise ValueError("Input tensor metadata not found for leaky_relu") + dtype = x_meta.dtype + + zero_slot = emit_lifted_constant(P, 0.0, dtype) + slope_slot = emit_lifted_constant(P, negative_slope, dtype) + + _, cond_slot = P.make_tmp_slot() + P.emit( + GreaterEqualNode( + a=P.slot_to_tid(x), + b=P.slot_to_tid(zero_slot), + out=P.slot_to_tid(cond_slot), + ) + ) + + _, scaled_slot = P.make_tmp_slot() + P.emit( + MultiplyNode( + a=P.slot_to_tid(slope_slot), + b=P.slot_to_tid(x), + out=P.slot_to_tid(scaled_slot), + ) + ) + + out = P.make_or_get_slot(n) + P.emit( + WhereNode( + condition=P.slot_to_tid(cond_slot), + x=P.slot_to_tid(x), + y=P.slot_to_tid(scaled_slot), + out=P.slot_to_tid(out), + ) + ) + return out + + @REGISTRY.register(target=[torch.ops.aten._log_softmax.default]) def _log_softmax_handler(P: MLXProgramBuilder, n: Node) -> Slot: """Handle aten._log_softmax.default - log of softmax. diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 8f52116f6b8..e96c8075903 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -405,6 +405,60 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (x,) +class LeakyReLUModel(nn.Module): + """Model that applies leaky_relu with an optional negative slope.""" + + def __init__(self, negative_slope: Optional[float] = 0.01): + super().__init__() + self.negative_slope = negative_slope + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.negative_slope is None: + return torch.nn.functional.leaky_relu(x) + return torch.nn.functional.leaky_relu(x, negative_slope=self.negative_slope) + + +@register_test +class LeakyReLUTest(OpTestCase): + """Test case for leaky_relu activation with various negative slopes.""" + + name = "leaky_relu" + rtol = 1e-5 + atol = 1e-5 + + def __init__( + self, + shape: Tuple[int, ...] = (2, 3, 4), + negative_slope: Optional[float] = 0.01, + ): + self.shape = shape + self.negative_slope = negative_slope + shape_str = "x".join(str(s) for s in shape) + slope_str = "default" if negative_slope is None else f"slope{negative_slope}" + self.name = f"leaky_relu_{slope_str}_{shape_str}" + + @classmethod + def get_test_configs(cls) -> List["LeakyReLUTest"]: + return [ + cls(shape=(2, 3, 4), negative_slope=0.01), + cls(shape=(2, 3, 4), negative_slope=None), + cls(shape=(4, 8), negative_slope=0.1), + cls(shape=(10,), negative_slope=0.2), + cls(shape=(10,), negative_slope=1.5), + cls(shape=(2, 8, 16), negative_slope=0.01), + ] + + def create_model(self) -> nn.Module: + return LeakyReLUModel(self.negative_slope) + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + numel = 1 + for size in self.shape: + numel *= size + x = torch.linspace(-4.0, 4.0, steps=numel).reshape(self.shape) + return (x,) + + class GELUModel(nn.Module): """Simple model using GELU activation."""