Function f

axlearn/common/optimizers.py:609–610 · view source on GitHub ↗

(g, p, s)

Source from the content-addressed store, hash-verified

607	param_scales = _weight_decay_scales(params, per_param_scale=per_param_scale)
608
609	def f(g, p, s):
610	return g + weight_decay * lr_scale * p.value * s
611
612	updates = jax.tree.map(
613	lambda x, y, z: None if x is None else f(x, y, z),

gMethod · 0.85

test_wrap_functionMethod · 0.85

test_with_pjit_out_shardingsMethod · 0.85

test_fused_grouped_qkv_inside_shard_mapMethod · 0.85

test_annotate_stackMethod · 0.85

test_walk_annotated_tbMethod · 0.85

test_in_context_exceptionMethod · 0.85

update_fnFunction · 0.85

test_init_states_kv_partition_specMethod · 0.85

runFunction · 0.85

grad_fnMethod · 0.85

no outgoing calls

gMethod · 0.68

test_wrap_functionMethod · 0.68

test_with_pjit_out_shardingsMethod · 0.68

test_fused_grouped_qkv_inside_shard_mapMethod · 0.68

test_annotate_stackMethod · 0.68

test_walk_annotated_tbMethod · 0.68

test_in_context_exceptionMethod · 0.68

test_init_states_kv_partition_specMethod · 0.68

grad_fnMethod · 0.68