zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-03 01:06:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add losses to the docs, fix black failur (#92)

Browse Source
This commit is contained in:
Awni Hannun 2023-12-09 06:06:52 -08:00 committed by GitHub
parent 430bfb4944
commit 2520dbcf0a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
docs/src/python/nn.rst
View File

@ -170,3 +170,13 @@ simple functions.
gelu_fast_approx gelu_fast_approx
relu relu
silu silu
Loss Functions
--------------
.. autosummary::
:toctree: _autosummary_functions
:template: nn-module-template.rst
losses.cross_entropy
losses.l1_loss

9
python/mlx/nn/losses.py
View File

@ -12,12 +12,9 @@ def cross_entropy(
Args: Args:
logits (mx.array): The predicted logits. logits (mx.array): The predicted logits.
targets (mx.array): The target values. targets (mx.array): The target values.
axis (int, optional): The axis over which to compute softmax. Defaults to -1. axis (int, optional): The axis over which to compute softmax. Default: ``-1``.
reduction (str, optional): Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. reduction (str, optional): Specifies the reduction to apply to the output:
'none': no reduction will be applied. ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
'mean': the sum of the output will be divided by the number of elements in the output.
'sum': the output will be summed.
Defaults to 'none'.
Returns: Returns:
mx.array: The computed cross entropy loss. mx.array: The computed cross entropy loss.