ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update cross_entropy_loss.py

Browse source
This commit is contained in:
Daniel Han 2024-11-03 21:38:19 -08:00
commit e23ef99c03
1 changed files with 6 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/kernels/cross_entropy_loss.py
View file

@ -38,9 +38,9 @@ def _cross_entropy_forward(
labels_ptr ,
VOCAB_SIZE : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
DO_LOGIT_SCALING ,
LOGIT_SCALE ,
):
"""
@ -112,9 +112,9 @@ def _chunked_cross_entropy_forward(
VOCAB_SIZE : tl.constexpr(tl.int32),
N_CHUNKS : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
DO_LOGIT_SCALING ,
LOGIT_SCALE ,
):
"""
@ -195,9 +195,9 @@ def _cross_entropy_backward(
labels_ptr ,
VOCAB_SIZE : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
DO_LOGIT_SCALING ,
LOGIT_SCALE ,
):
"""