stuff

Author: jlamypoirier

fast_llm/functional/cross_entropy.py fast_llm/functional/entropy_loss.pyfast_llm/functional/cross_entropy.py renamed to fast_llm/functional/entropy_loss.py

@@ -77,7 +77,7 @@ def get_loss(
         sequence_parallel_logits: bool = False,
         kwargs: dict[str, typing.Any],
     ) -> "tuple[torch.Tensor, torch.Tensor | None]":
-        from fast_llm.functional.cross_entropy import entropy_loss_forward_backward
+        from fast_llm.functional.entropy_loss import entropy_loss_forward_backward
 
         labels = kwargs[LanguageModelKwargs.labels]
 
@@ -167,7 +167,7 @@ def get_loss(
         sequence_parallel_logits: bool = False,
         kwargs: dict[str, typing.Any],
     ) -> "tuple[torch.Tensor, torch.Tensor | None]":
-        from fast_llm.functional.cross_entropy import entropy_loss_forward_backward
+        from fast_llm.functional.entropy_loss import entropy_loss_forward_backward
 
         if prediction_distance > 0:
             raise NotImplementedError()

fast_llm/layers/language_model/loss/config.py

@@ -5,7 +5,7 @@
 
 from fast_llm.engine.distributed.config import DistributedBackend
 from fast_llm.functional.config import EntropyLossImplementation, EntropyLossType, TargetFormat, TritonConfig
-from fast_llm.functional.cross_entropy import entropy_loss_forward_backward
+from fast_llm.functional.entropy_loss import entropy_loss_forward_backward
 from fast_llm.utils import Assert
 from tests.utils.subtest import DistributedTestContext