Fix step activation docstring to match >= threshold behavior

[ayaangazali]

Proposed changes

The step function and Step module docstrings say the output is "1 if the input is greater than a specified threshold", but the implementation is mx.where(x >= threshold, 1, 0) and the math block right below the prose already uses x >= threshold. So the prose was the only part that disagreed. Quick check on main:

import mlx.core as mx
import mlx.nn as nn

nn.Step(0.0)(mx.array([-1.0, 0.0, 1.0])).tolist()  # [0, 1, 1]  -> 0 maps to 1, i.e. >=

Updated the prose to "greater than or equal to" in both places, and documented the threshold default (0.0) while I was there.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

(docstring-only, test_nn.py still passes 69/69)

---

quick note: i'm a freshman still learning, and i lean on Claude Code to surface these doc/code mismatches, but i ran the snippet above myself to confirm the boundary value maps to 1 before writing this. happy to adjust if you'd word it differently.

[ayaangazali]

Gentle nudge on three small docstring PRs that have been sitting: this one, #3903 and #3906. Each is a one or two line correction where the text disagrees with the code right next to it:

• Fix step activation docstring to match >= threshold behavior #3902 (this one): step and Step say the output is 1 if the input is "greater than" the threshold, but the code is mx.where(x >= threshold, 1, 0) and the math block just below the prose already uses >=.
• Fix InstanceNorm Shape docstring to require at least 3 dimensions #3903: InstanceNorm documents its input as (..., C), but __call__ raises for anything with fewer than 3 dimensions.
• Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn #3906: filter_and_map describes the callback arguments as (value, key, module), while the call site is filter_fn(self, k, v), which matches the type hints.

I rebased all three on current main just now and they are green. No rush from my side. Happy to squash them into one PR if that is easier to review, or to close any of them if you would rather not take the change.