Step

Applies the Step Activation Function (Heaviside).

Abstract Signature:

Step(threshold: float = 0.0)

PyTorchApple MLXFlax NNX

PyTorch

API: —

Strategy: Plugin (torch_step_layer)

Apple MLX

API: mlx.nn.Step

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: nnx.training.optimizer.Optimizer.step

Strategy: Direct Mapping

Official Docs ↗