Skip to content

Rewrite of training, callbacks, losses, optimizers and batchers for more flexibility. #72

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jaosch wants to merge 9 commits into
base: develop
Choose a base branch
from
Draft

Rewrite of training, callbacks, losses, optimizers and batchers for more flexibility. #72

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
0b23910
Modularized fit
jaosch Sep 12, 2025
75a3f96
Concretization of type annotations + example for `fit_core`
jaosch Sep 12, 2025
08e4f6b
Tiny fix
jaosch Sep 12, 2025
167fc1c
Enable stateful batching + training state extended with treedefs
jaosch Sep 13, 2025
a7fc3a2
Added EvaluationContext and TrainingContext
jaosch Sep 13, 2025
5294cc7
Added call of on_training_end
jaosch Sep 13, 2025
66dcc6c
Pass full `TrainingContext` into the batcher
jaosch Sep 13, 2025
e35d89a
New HistoryCallback + distribution over multiple files
jaosch Sep 13, 2025
4c2b930
Set batch_data back to its original functionality
jaosch Sep 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 7 additions & 10 deletions klax/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,15 @@
from ._callbacks import (
Callback as Callback,
)
from ._callbacks import (
CallbackArgs as CallbackArgs,
)
from ._callbacks import (
HistoryCallback as HistoryCallback,
)
from ._context import (
EvaluationContext,
TimingInfo,
TrainingContext,
TrainingState,
)
from ._datahandler import (
BatchGenerator as BatchGenerator,
)
Expand All @@ -32,13 +35,7 @@
split_data as split_data,
)
from ._losses import (
MAE as MAE,
)
from ._losses import (
MSE as MSE,
)
from ._losses import (
Loss as Loss,
LossFactory as LossFactory,
)
from ._losses import (
mae as mae,
Expand Down
Loading
Loading