File size: 1,114 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import kfac_jax
from .production import (
KFACBundle,
apply_finetune_kfac_step,
apply_router_kfac_step,
init_finetune_kfac_state,
init_router_kfac_state,
)
from .targets import process_finetune_targets, process_route_targets
def learning_rate(config, step: int) -> float:
return float(config.learning_rate_numerator) / (
float(config.learning_rate_offset)
+ float(step) / float(config.learning_rate_decay_steps)
)
def register_scale_and_shift(y, x, scale, tag_id: str):
from hamiltonzero.model.tree import _kfac_name_kw
return kfac_jax.register_scale_and_shift(
y,
x,
scale=scale,
shift=None,
**_kfac_name_kw(tag_id),
)
__all__ = [
"KFACBundle",
"apply_finetune_kfac_step",
"apply_router_kfac_step",
"init_finetune_kfac_state",
"init_router_kfac_state",
"learning_rate",
"process_finetune_targets",
"process_route_targets",
"register_scale_and_shift",
]
|