Function update_fn

axlearn/common/optimizers.py:99–105

(
        updates: NestedTensor, state: dict[str, Any], params: NestedOptParam
    )

        return {k: v.init(params) for k, v in kwargs.items()}

    def update_fn(
        updates: NestedTensor, state: dict[str, Any], params: NestedOptParam
    ) -> tuple[NestedTensor, optax.EmptyState]:
        new_state = {}
        for k, v in kwargs.items():
            updates, new_state[k] = v.update(updates, state[k], params)
        return updates, new_state

    def partition_fn(param_spec):
        return {k: v.partition(param_spec) for k, v in kwargs.items()}