Additoinal losses by JayYun98 · Pull Request #336 · ml-explore/mlx

JayYun98 · 2024-01-01T17:56:49Z

Proposed changes

Added commonly used losses with tests

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)

awni

Thanks for the contribution!!

We already have a hinge and huber loss already. The focal, dice, and contrastive are in #324. Please remove all of them and rebase to resolve conflicts.

If you want to add the cosine loss in this PR that would be great.

add cos sim

cos sim test

JayYun98 · 2024-01-02T14:12:58Z

Hi @awni , rebased and make cos loss to remain on this pr.

Also commented on #324 based on my implementation.

Thanks :)

awni

Thanks! I left a couple comments, please address then let me know when ready for re-review.

awni · 2024-01-03T03:58:33Z


 if __name__ == "__main__":
    unittest.main()
+


Nit: can you remove that extra line

awni · 2024-01-03T14:58:59Z

Hey @Jyun1998 sorry but somehow my main comment did not get included. I must have not pushed the save button by accident.

Basically I'm wondering what reference you use for this loss as it looks quite different than the similarly named cosine similarity loss in PyTorch ?

For example there is a target and a margin both of which I do not expect. Unless there is a good reason for the difference, I would suggest we follow the PyTorch implementation as a reference.

awni · 2024-01-05T03:59:00Z

@Jyun1998 are you still planning to follow up on this?

JayYun98 · 2024-01-07T14:15:52Z

@Jyun1998 are you still planning to follow up on this?

https://github.com/keras-team/keras/blob/v2.14.0/keras/losses.py#L1162-L1236

Hi awni, according to common tensorflow and pytorch implementaiton, the functions are composed of doing l2 norm to each embedding and returns the negative of dot product of both embedding.

My codes also do so, and margin-based loss is applying if it's necessary :)

awni · 2024-01-07T14:18:55Z

@Jyun1998 got it. We should keep it simple until we see that we need more features. Could you follow the PyTorch cosine similarity loss? I think that one covers the most common case (margin / targets are are niche and possibly never needed so we don't want to add them to the API until we are sure they are necessary)

JayYun98 · 2024-01-07T14:27:19Z

@Jyun1998 got it. We should keep it simple until we see that we need more features. Could you follow the PyTorch cosine similarity loss? I think that one covers the most common case (margin / targets are are niche and possibly never needed so we don't want to add them to the API until we are sure they are necessary)

I also agree. Even though there's also margin for pytorch F.cosine_similarity, it defaultly do not use it.

There's only slight change needed for changes and I will test the function and commit asap. Thanks for the review

awni · 2024-01-07T14:30:49Z

Even though there's also margin for pytorch F.cosine_similarity

I don't see the margin in the docs? Is it in the source code?

JayYun98 · 2024-01-07T14:35:47Z

Even though there's also margin for pytorch F.cosine_similarity

I don't see the margin in the docs? Is it in the source code?

https://pytorch.org/docs/stable/generated/torch.nn.CosineEmbeddingLoss.html#torch.nn.CosineEmbeddingLoss

awni · 2024-01-07T14:38:47Z

I see, thanks. Yes let's go with the plain cosine similarity for now. Thank you!

awni · 2024-01-07T14:39:07Z

Also could you rebase and resolve conflicts?

JayYun98 · 2024-01-07T15:05:35Z

Also could you rebase and resolve conflicts?

Am I correct that losses test codes are gone?

nvm found new losses test file

* fix: tidy pull request template * fix: feedback from awni

* Added scale for RoPE * Ran pre-commit * Added RoPE scaling test * Added docstring for scale parameter * Modified docstrings

* feat: Add Dropout3d layer to nn.layers * acknowledgement * Add dropout tests to test_nn.py * run pre-commit * Add activation functions and dropout3d ops * Add dropout tests for bfloat16 and float16

Co-authored-by: xuan <xuan@apple.com>

Co-authored-by: Chunyang Wen <chunyang_wen@apple.com>

* spelling: accumulates Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: across Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: additional Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: against Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: among Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: array Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: at least Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: available Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: axes Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: basically Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: bfloat Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: bounds Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: broadcast Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: buffer Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: class Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: coefficients Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: collision Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: combinations Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: committing Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: computation Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: consider Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: constructing Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: conversions Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: correctly Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: corresponding Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: declaration Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: default Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: dependency Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: destination Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: destructor Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: dimensions Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: divided Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: element-wise Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: elements Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: endianness Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: equivalent Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: explicitly Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: github Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: indices Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: irregularly Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: memory Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: metallib Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: negative Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: notable Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: optional Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: otherwise Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: overridden Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: partially Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: partition Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: perform Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: perturbations Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: positively Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: primitive Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: repeat Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: repeats Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: respect Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: respectively Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: result Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: rounding Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: separate Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: skipping Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: structure Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: the Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: transpose Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: unnecessary Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: unneeded Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> * spelling: unsupported Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com> --------- Signed-off-by: Josh Soref <2119212+jsoref@users.noreply.github.com>

add cos sim

cos sim test

JayYun98 · 2024-01-07T15:59:10Z

Also could you rebase and resolve conflicts?

Could you check? Thanks :)

awni

A few more changes please.

awni

Looks great, thanks for the contribution @Jyun1998 !

feat: `ARRAY_API_TESTS_MODULE` for runtime-defined xp

junwoo-yun and others added 6 commits December 31, 2023 23:20

change class order

f935cfb

add losses

bc2fa2a

change after all test passes

80b4d84

change after all tests

19bc6a3

precommit

80c4630

black

353a4cb

awni requested changes Jan 1, 2024

View reviewed changes

JayYun98 mentioned this pull request Jan 2, 2024

feat: Add new loss functions for neural networks #324

Closed

4 tasks done

JayYun98 added 3 commits January 2, 2024 23:04

Merge branch 'main' into additoinal_losses

6de4c7a

Update losses.py

d66fec6

add cos sim

Update test_nn.py

d408541

cos sim test

awni requested changes Jan 3, 2024

View reviewed changes

Update test_nn.py

975d652

JayYun98 added 3 commits January 7, 2024 23:58

change after review

14157f3

nn losses

72299db

change afte rewivew

acf37ac

JayYun98 and others added 3 commits January 8, 2024 00:24

rebase

35a1b9e

Docs update (ml-explore#144)

3490e69

fix: tidy pull request template (ml-explore#143)

9f0be58

* fix: tidy pull request template * fix: feedback from awni

angeloskath and others added 16 commits January 8, 2024 00:28

Improve repeat using broadcasting and reshape (ml-explore#318)

9eac811

Added implementation for Scaled RoPE. (ml-explore#261)

e8128ab

* Added scale for RoPE * Ran pre-commit * Added RoPE scaling test * Added docstring for scale parameter * Modified docstrings

fix typos (ml-explore#327)

021268d

feat: Add Dropout3d layer to nn.layers (ml-explore#313)

62aede2

* feat: Add Dropout3d layer to nn.layers * acknowledgement * Add dropout tests to test_nn.py * run pre-commit * Add activation functions and dropout3d ops * Add dropout tests for bfloat16 and float16

fix type cast error in item() for bfloat16 (ml-explore#339)

4e60305

Co-authored-by: xuan <xuan@apple.com>

Remove useless import (ml-explore#340)

ffa2d26

Co-authored-by: Chunyang Wen <chunyang_wen@apple.com>

Update losses.py

45a0ae7

add cos sim

Update test_nn.py

4a3599a

cos sim test

change afte rewivew

7288672

Rebase

d03ef7c

rebase

9ab65ca

rebase

01a5625

rebase

46e12c3

rebase

2d371c9

Merge branch 'main' into additoinal_losses

e291fd1

JayYun98 added 3 commits January 8, 2024 01:00

rebasE

c2a61f1

space between docstring

f38c541

run black

5127475

awni requested changes Jan 7, 2024

View reviewed changes

JayYun98 and others added 5 commits January 8, 2024 01:45

fix by comment

d17b4d5

change by review

5991471

change by review

3d8871a

- to + for test cases

469c400

nits in docstring

7778b9c

awni approved these changes Jan 8, 2024

View reviewed changes

awni merged commit 0b8aedd into ml-explore:main Jan 8, 2024

awni mentioned this pull request Jan 23, 2024

Add binary focal cross entropy loss function #498

Closed

4 tasks done

zcbenz pushed a commit to Ritabanm/mlx that referenced this pull request Sep 30, 2026

Merge pull request ml-explore#336 from lucascolley/runtime-xp

a882502

feat: `ARRAY_API_TESTS_MODULE` for runtime-defined xp

Conversation

JayYun98 commented Jan 1, 2024

Proposed changes

Checklist

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

JayYun98 commented Jan 2, 2024

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

awni Jan 3, 2024

Choose a reason for hiding this comment

Uh oh!

JayYun98 Jan 3, 2024

Choose a reason for hiding this comment

Uh oh!

awni commented Jan 3, 2024

Uh oh!

awni commented Jan 5, 2024

Uh oh!

JayYun98 commented Jan 7, 2024

Uh oh!

awni commented Jan 7, 2024

Uh oh!

JayYun98 commented Jan 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

awni commented Jan 7, 2024

Uh oh!

JayYun98 commented Jan 7, 2024

Uh oh!

awni commented Jan 7, 2024

Uh oh!

awni commented Jan 7, 2024

Uh oh!

JayYun98 commented Jan 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JayYun98 commented Jan 7, 2024

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants

JayYun98 commented Jan 7, 2024 •

edited

Loading

JayYun98 commented Jan 7, 2024 •

edited

Loading