Skip to content

Fix log_cosh_loss docstring to document the element-wise loss#3846

Open
winklemad wants to merge 1 commit into
ml-explore:mainfrom
winklemad:fix-log-cosh-docstring
Open

Fix log_cosh_loss docstring to document the element-wise loss#3846
winklemad wants to merge 1 commit into
ml-explore:mainfrom
winklemad:fix-log-cosh-docstring

Conversation

@winklemad

Copy link
Copy Markdown

Proposed changes

log_cosh_loss's docstring documents the loss as a mean-reduced scalar:

logcosh(y_true, y_pred) = (1/n) Σ log(cosh(y_pred - y_true))

but the function defaults to reduction="none" and returns the per-element log(cosh(y_pred - y_true)). The (1/n) Σ only corresponds to reduction="mean", so at the default the documented formula doesn't match the output:

>>> import mlx.core as mx, mlx.nn as nn
>>> nn.losses.log_cosh_loss(mx.array([1.0, 2.0]), mx.array([0.0, 0.0]))
array([0.433781, 1.325], dtype=float32)   # per-element, not the scalar 0.87939 the formula implies

It's also the only loss in losses.py whose math block bakes in the reduction — huber_loss, smooth_l1_loss, hinge_loss, and the others all document the element-wise loss and delegate aggregation to the reduction argument. This drops the (1/n) Σ so the formula reads the element-wise log(cosh(y_pred - y_true)), matching both the code and the rest of the file.

Docs-only; no behavior change.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit (black + isort) on the changed file
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@Pablosinyores Pablosinyores left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed against the source: loss = mx.logaddexp(errors, -errors) - math.log(2) is log(cosh(errors)) computed element-wise, and reduction defaults to none. The old (1/n) sum form baked in a mean that the default reduction doesn't apply, so it contradicted the actual output. The element-wise form now matches the code and is consistent with how huber_loss above documents its formula (reduction left to the Args). LGTM.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants