Method update

axlearn/common/learner.py:245–264

Computes `model_params` updates from `update`. Args: updates: The updates to potentially transform and then apply. Returns: The updated model parameters. The learner state updates will be placed in the output collection's 'state_update' section.

(self, updates: Updates)

        cfg = self.config
        return jax.tree.map(
            lambda path: match_regex_rules(
                path, rules=cfg.update_rules, default_value=UpdateType.ALL_UPDATES
            ),
            tree_paths(tree),
        )

    def should_update_with_optimizers(self, model_params: Nested[OptParam]) -> dict:
        """Returns whether each parameter should be updated with the optimizers.

        Args:
            model_params: A nested structure with OptParams as leaf nodes.

        Returns:
            A nested dict with the same structure as `model_params` with boolean leaf values.
        """
        return jax.tree.map(should_update_with_optimizers, self._update_types(model_params))

    def handles_undefined_loss(self) -> bool:
        return self.config.ignore_undefined_loss

    def update(self, updates: Updates) -> Nested[Tensor]:
        """Computes `model_params` updates from `update`.