Prompt:
 if self._count % self.gas == 0:
            self.optimizer_step()

            # param callback (e.g., parameter clipping)
            if self.is_implemented("param_callback"):
                self.param_callback()

            if self._strategy != "default" and self._count % (self.gas * 20) == 0:
                self.synchronize_params(self.trainable_parameters())

            # zero-out grad
            self.zero_grad()

        return loss_dict

    def step_normal(self, global_step=None):
        if self.check_ready():
            # loop start
            if self._inner_loop_start:
                if self.is_implemented("on_inner_loop_start"):
                    self.on_inner_loop_start()
                self._inner_loop_start = False

                # copy current parameters, buffers, optimizer states
                if self._roll_back:
                    self.cache_states()

            # increase count (local step)
            if self._training:
                self._count += 1

            # one step grdient descent
            loss_dict = self.one_step_descent()

            # lr scheduler step
            if self.scheduler is not None and not self._roll_back:
                self.scheduler.step()

            # logging
            if (
                self.log_step > 0
                and self._count % self.log_step == 0
                and self.is_rank_zero()
            ):
                self.log(loss_dict, global_step)

            # call parent step_normal after unrolling
            if (
                self._training
                and self._count % (self._unroll_steps * self.gas) == 0
                and self._count > self.warmup_steps
            ):
                for problem in self._parents:
                    idx = problem.children.index(self)
                    problem.ready[idx] = True
                    problem.step_normal(global_step=global_step)

                self._inner_loop_start = True

            self.ready = [False for _ in range(len(self._children))]

    def step_after_roll_back(self):
        if self.check_ready() and self._training:
            if self._roll_back:
                # recover from cached states
                self.recover_states()

                # one step gradient step
                _ = self.one_step_descent(batch=self.cur_batch)

                # lr scheduler
                if self.scheduler is not None:
                    self.scheduler.step()

                # call parent step_after_roll_back
                for problem in self._parents:
                    idx = problem.children.index(self)
                    problem.ready[idx] = True
                    problem.step_after_roll_back()

            self.ready = [False for _ in range(len(self._children))]

    def step(self, global_step=None):
        """
        ``step`` method abstracts a one-step gradient descent update with four sub-steps:
        1) data loading, 2) cost calculation, 3) gradient calculation, and 4) parameter update.
        It also calls upper-level problems' step methods after unrolling gradient steps based on
        the hierarchical dependency graph.

        :param global_step: global step of the whole multilevel optimization. Defaults to None.
        :type global_step: int, optional
        """
        self._global_step = global_step
        self.step_normal(global_step=global_step)
        if (
            self._count % (self._unroll_steps * self.gas) == 0
            and self._count > self.warmup_steps
        ):
            self.step_after_roll_back()

    def get_batch(self):
        """
        Load training batch from the user-provided data loader

        :return: New training batch
        :rtype: Any
        """
        batch = tuple(
            self.get_batch_single_loader(i) for i in range(len(self.train_data_loader))
        )

        return batch[0] if len(batch) == 1 else batch

    def get_batch_single_loader(self, idx):
        """
        Load training batch from one of the user-provided data loader(s)

        :return: New training batch
        :rtype: Any
        """
        data_iterator = self.train_data_iterator[idx]
        try:
            batch = next(data_iterator)
        except StopIteration:
            if idx == 0:
                self.epoch_callback_exec()
            self.epoch_counter[idx] += 1
            train_data_loader = self.train_data_loader[idx]
            if self._strategy in ["distributed", "zero", "fsdp"]:
                train_data_loader.set_epoch(self.epoch_counter[idx])
            self.train_data_iterator[idx] = iter(train_data_loader)
            batch = next(self.train_data_iterator[idx])
        if not isinstance(batch, dict):
            batch = tuple(
                convert_tensor(value, self.device, self._is_default_fp16())
                for value in batch
            )
        else:
            for key, value in batch.items():
                batch[key] = convert_tensor(value, self.device, self._is_default_fp16())

        return batch

    def get_loss(self, batch):
        """
        Calculate loss and log metrics for the current batch based on the user-defined loss
        function.

        :return: loss and log metrics (e.g. classification accuracy)
        :rtype: dict
        """
        maybe_loss_dict = self.training_step_exec(batch)
        is_dict = isinstance(maybe_loss_dict, dict)
        loss = maybe_loss_dict["loss"] if is_dict else maybe_loss_dict
        loss_no_scale = loss.item()
        if self._is_default_fp16():
            loss = self.scaler.scale(loss)
        loss = loss / self.gas

        # construct loss dict
        loss_dict = {"loss": loss_no_scale}
        if is_dict:
            for key, value in maybe_loss_dict.items():
                if key != "loss":
                    loss_dict[key] = value

        return loss, loss_dict

    def backward(
        self,
        loss,
        params,
        paths,
        create_graph=False,
        retain_graph=True,
        allow_unused=True,
    ):
        """
        Calculate the gradient of ``loss`` with respect to ``params`` based on a user-defined
        ``config``.

        :param loss: Outputs of the differentiated function.
        :type loss: Tensor
        :param params: Inputs with respect to which the gradient will be returned.
        :type params: Sequence of Tensor
        :param paths: Paths on which the gradient will be calculated.
        :type paths: List of list of Problem
        :param create_graph:
            If ``True``, graph of the derivative will be constructed, allowing to compute higher order
            derivative products. Default: ``True``.
        :type create_graph: bool, optional
        :param retain_graph:
            If ``False``, the graph used to compute the grad will be freed. Note that in nearly all
            cases setting this option to ``True`` is not needed and often can be worked around in a much
            more efficient way. Defaults to the value of ``create_graph``.
        :type retain_graph: bool, optional
        :param allow_unused:
            If ``False``, specifying inputs that were not used when computing outputs (and therefore
            their grad is always zero) is an error. Defaults to ``False``.
        :type allow_unused: bool, optional
        """
        # direct grad
        if len(paths) > 0 or not self.gradient_accumulation_boundary():
            grads = torch.autograd.grad(
                loss,
                params,
                create_graph=create_graph,
                retain_graph=retain_graph,
                allow_unused=allow_unused,
            )
            self.set_grads(params, grads)
        else:
            torch.autograd.backward(
                loss,
                inputs=params,
                create_graph=create_graph,
                retain_graph=retain_graph,
            )

        # indirect grad: best-response Jacobian
        if self._config.first_order:
            for idx, path in enumerate(paths):
                retain_graph_implicit = False if idx == len(paths) - 1 else True
                do_sync = bool(
                    idx == len(paths) - 1 and self.gradient_accumulation_boundary()
                )
                grads = get_grads(loss, path, retain_graph_implicit, do_sync)
                if not do_sync:
                    self.set_grads(params, grads)

    def set_grads(self, params, grads):
        """
        Set gradients for trainable parameters. ``params.grad = grads``

        :param params: Trainable parameters
        :type params: Sequence of Tensor
        :param grads: Calculated gradient
        :type grads: Sequence of Tensor
        """
        for param, grad in zip(params, grads):
            if grad is not None:
                if hasattr(param, "grad") and param.grad is not None:
                    param.grad = param.grad + grad
                else:
                    param.grad = grad

    def synchronize_params(self, params):
        """
        synchronize parameters across distributed data-parallel processes
        """
        if self._world_size > 1 and self._strategy not in ["fsdp", "accelerate"]:
            for param in params:
                dist.broadcast(param.data, 0)

    @abc.abstractmethod
    def optimizer_step(self, *args, **kwargs):
        """
        Update weights as in PyTorch's native ``optim.step()``
        """
        raise NotImplementedError

    def zero_grad(self):
        """
        Set gradients for trainable parameters for the current problem to 0.
        Similar with PyTorch's ``optim.zero_grad()`` or ``module.zero_grad()``.
        """
        for param in list(self.trainable_parameters()):
            if hasattr(param, "grad"):
                del param.grad

    def clip_grad(self):
        """
        Perform gradient clipping based on the norm provided by Config
        """
        if self._strategy != "fsdp":
            torch.nn.utils.clip_grad_norm_(
                parameters=self.trainable_parameters(), max_norm=self.gradient_clipping
            )
        else:
            self.module.clip_grad_norm_(max_norm=self.gradient_clipping)

    def state_dict(self):
        """
        Return all states involved in ``Problem`` with a Python dictionary. By default, it
        includes ``self.module.state_dict`` and ``self.optimizer.state_dict``. Depending on users'
        configurations, it may include ``self.scheuler.state_dict`` (lr scheduler) and
        ``self.scaler.state_dict`` (fp16 training)
        """
        state_dict = {}
        state_dict["module"] = self.module.state_dict()
        state_dict["optimizer"] = self.optimizer.state_dict()
        if self.scheduler is not None:
            state_dict["scheduler"] = self.scheduler.state_dict()
        if self._is_default_fp16():
            state_dict["scaler"] = self.scaler.state_dict()

        return state_dict

    def load_state_dict(self, state_dict):
        """Load the state for the ``Problem``

        Args:
            state_dict (dict): Python dictionary of Problem states.
        """
        self.module.load_state_dict(state_dict["module"])
        self.optimizer.load_state_dict(state_dict["optimizer"])
        if self.scheduler is not None and "scheduler" in state_dict:
            self.scheduler.load_state_dict(state_dict["scheduler"])
        if self._is_default_fp16() and "scaler" in state_dict:
            self.scaler.load_state_dict(state_dict["scaler"])

    def configure_distributed_training(self, dictionary):
        """
        Set the configuration for distributed training.

        :param dictionary: Python dictionary of distributed training provided by Engine.
        :type dictionary: dict
        """
        self._strategy = dictionary["strategy"]
        self._backend = dictionary["backend"]
        self._world_size = dictionary["world_size"]
        self._rank = dictionary["rank"]
        self._local_rank = dictionary["local_rank"]

    def configure_roll_back(self, roll_back):
        """
        Set the roll-back (warm- start) option from Engine

        :param roll_back: roll-back (warm-start) on/off
        :type roll_back: bool
        """
        if len(self._parents) > 0:
            self._roll_back = roll_back

    def configure_device(self, device):
        """
        Set the device for the current problem.
        """
        self.device = device

    def get_opt_param_group_for_param(self, param):
        """
        Get optimizer param_group for specific parameter

        :param param: Parameter for which optimizer param_group is inquired
        :type param: torch.nn.Parameter
        :return: param_group for the given parameter
        :rtype: dict
        """
        param_groups = self.optimizer.param_groups
        for group in param_groups:
            for p in group["params"]:
                if param is p:
                    return group

    def get_opt_state_for_param(self, param):
        """
        Get optimizer state for specific parameter

        :param param: Parameter for which optimizer state is inquired
        :type param: torch.nn.Parameter
        :return: optimizer state for the given parameter
        :rtype: dict
        """
        state = self.optimizer.state
        return state[param]

    @abc.abstractmethod
    def cache_states(self):
        """
        Cache params, buffers, optimizer states when ``config.roll_back`` is set to ``True`` in
        ``step``.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def recover_states(self):
        """
        Recover params, buffers, optimizer states when ``config.roll_back`` is set to ``True`` in
        ``step``.
        """
        raise NotImplementedError

    def epoch_callback_exec(self):
        if self.is_implemented("epoch_callback"):
            self.epoch_callback()

    def gradient_accumulation_boundary(self):
        """
        Check whether the current step is on the gradient accumulation boundary
        """
        return bool(self._count % self.gas == 0)

    def _is_default_fp16(self):
        """
        Check whether to use PyTorch native fp16 (mixed-precision) feature
        """
        if not self._fp16 or self._strategy in ["accelerate"]:
            return False
        return True

    def is_implemented(self, fn_name):
        """
        Check if ``fn_name`` method is implemented in the class

        :rtype: bool
        """
        return callable(getattr(self, fn_name, None))

    def check_ready(self):
        """
        Check if unrolling processes of lower level problems in the hierarchical dependency
        graph are all ready/done. ``step`` function is only excuted when this method returns
        ``True``.

        :rtype: bool
        """
        return all(self.ready)

    def log(self, stats, global_step):
        """
        Log (training) stats to the ``self.logger``

        :param stats: log metrics such as loss and classification accuracy.
        :type stats: Any
        :param step: global/local step associated with the ``stats``.
        :type step: int
        """

        loss_log = log_from_loss_dict(stats)
        if global_step is None:
            self.logger.info(
                f'[Problem "{self._name}"] [Local Step {self._count}] {loss_log}'
            )
        else:
            self.logger.info(
                f'[Problem "{self._name}"] [Global Step {global_step}] [Local Step {self._count}] '
                f"{loss_log}"
            )
        cur_step = global_step
        if global_step is None or self.log_local_step:
            cur_step = self._count
        self.logger.log(stats, tag=self._name, step=cur_step)


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from betty.problems import Problem

# pylint: disable=W0223
class ImplicitProblem(Problem):
    """
    ``ImplicitProblem`` is sublassed from ``Problem``.
    """

    def __init__(
        self,
        name,
        config,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):

        super().__init__(
            name,
            config,
            module,
            optimizer,
            scheduler,
            train_data_loader,
            extra_config,
        )
        self.module_state_dict_cache = None
        self.opitmizer_state_dict_cache = None


Prompt:
import torch

def convert_tensor(item, device=None, fp16=False):

    if not isinstance(item, torch.Tensor):
        return item
    return item.to(device)


Prompt:
import torch

def convert_tensor(item, device=None, fp16=False):
    if not isinstance(item, torch.Tensor):
        return item
    return item.to(device)

def get_grad_norm(parameters):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))

    total_norm = 0.0
    for p in parameters:
        param_norm = p.grad.data.float().norm()
        total_norm += param_norm.item() ** 2

    if (
        total_norm == float("inf")
        or total_norm == -float("inf")
        or total_norm != total_norm
    ):
        total_norm = -1

    return total_norm

def get_weight_norm(parameters):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    total_norm = 0.0
    for p in parameters:
        param_norm = torch.norm(p, dtype=torch.float32)
        total_norm += param_norm.item() ** 2

    if (
        total_norm == float("inf")
        or total_norm == -float("inf")
        or total_norm != total_norm
    ):
        total_norm = -1

    return total_norm

def flatten_list(regular_list):
    """[summary]
    Flatten list of lists
    """
    if type(regular_list[0] == list):
        return [item for sublist in regular_list for item in sublist]
    return regular_list

def get_param_index(param, param_list):
    param_list = list(param_list)
    for idx, p in enumerate(param_list):
        if p is param:
            return idx
    print("no corresponding parameter found!")

def get_multiplier(problem):
    if problem.leaf:
        return 1

    assert len(problem.children) > 0
    # stack to store all the nodes of tree
    s1 = []
    # stack to store all the leaf nodes
    s2 = []

    s1.append((problem, 1))
    while len(s1) != 0:
        curr, multiplier = s1.pop(0)

        if len(curr.children) != 0:
            for child in curr.children:
                s1.append((child, multiplier * curr.config.step))
        else:
            s2.append(multiplier)

    assert all(x == s2[0] for x in s2)
    return s2[0]

def log_from_loss_dict(loss_dict):

    outputs = []
    for key, values in loss_dict.items():
        if isinstance(values, dict) or isinstance(values, list):
            for value_idx, value in enumerate(values):
                full_key = key + "_" + str(value_idx)
                if torch.is_tensor(value):
                    value = value.item()
                output = f"{full_key}: {value}"
                outputs.append(output)
        else:
            if torch.is_tensor(values):
                values = values.item()
            output = f"{key}: {values}"
            outputs.append(output)
    return " || ".join(outputs)


Prompt:
import torch

def convert_tensor(item, device=None, fp16=False):
    if not isinstance(item, torch.Tensor):
        return item
    return item.to(device)

def get_grad_norm(parameters):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))

    total_norm = 0.0
    for p in parameters:
        param_norm = p.grad.data.float().norm()
        total_norm += param_norm.item() ** 2

    if (
        total_norm == float("inf")
        or total_norm == -float("inf")
        or total_norm != total_norm
    ):
        total_norm = -1

    return total_norm

def get_weight_norm(parameters):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    total_norm = 0.0
    for p in parameters:
        param_norm = torch.norm(p, dtype=torch.float32)
        total_norm += param_norm.item() ** 2

    if (
        total_norm == float("inf")
        or total_norm == -float("inf")
        or total_norm != total_norm
    ):
        total_norm = -1

    return total_norm

def flatten_list(regular_list):
    """[summary]
    Flatten list of lists
    """
    if type(regular_list[0] == list):
        return [item for sublist in regular_list for item in sublist]
    return regular_list

def get_param_index(param, param_list):
    param_list = list(param_list)
    for idx, p in enumerate(param_list):
        if p is param:
            return idx
    print("no corresponding parameter found!")

def get_multiplier(problem):
    if problem.leaf:
        return 1

    assert len(problem.children) > 0
    # stack to store all the nodes of tree
    s1 = []
    # stack to store all the leaf nodes
    s2 = []

    s1.append((problem, 1))
    while len(s1) != 0:
        curr, multiplier = s1.pop(0)

        if len(curr.children) != 0:
            for child in curr.children:
                s1.append((child, multiplier * curr.config.step))
        else:
            s2.append(multiplier)

    assert all(x == s2[0] for x in s2)
    return s2[0]

def log_from_loss_dict(loss_dict):
    outputs = []
    for key, values in loss_dict.items():
        if isinstance(values, dict) or isinstance(values, list):
            for value_idx, value in enumerate(values):
                full_key = key + "_" + str(value_idx)
                if torch.is_tensor(value):
                    value = value.item()
                output = f"{full_key}: {value}"
                outputs.append(output)
        else:
            if torch.is_tensor(values):
                values = values.item()
            output = f"{key}: {values}"
            outputs.append(output)
    return " || ".join(outputs)

def to_vec(tensor_list, alpha=1.0):
    return torch.cat([alpha * t.reshape(-1) for t in tensor_list])

def count_parameters(tensor_list):
    return sum([tensor.numel() for tensor in tensor_list])

def neg_with_none(a):

    if a is None:
        return None
    else:
        return -a


Prompt:
import torch

def convert_tensor(item, device=None, fp16=False):
    if not isinstance(item, torch.Tensor):
        return item
    return item.to(device)

def get_grad_norm(parameters):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))

    total_norm = 0.0
    for p in parameters:
        param_norm = p.grad.data.float().norm()
        total_norm += param_norm.item() ** 2

    if (
        total_norm == float("inf")
        or total_norm == -float("inf")
        or total_norm != total_norm
    ):
        total_norm = -1

    return total_norm

def get_weight_norm(parameters):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    total_norm = 0.0
    for p in parameters:
        param_norm = torch.norm(p, dtype=torch.float32)
        total_norm += param_norm.item() ** 2

    if (
        total_norm == float("inf")
        or total_norm == -float("inf")
        or total_norm != total_norm
    ):
        total_norm = -1

    return total_norm

def flatten_list(regular_list):
    """[summary]
    Flatten list of lists
    """
    if type(regular_list[0] == list):
        return [item for sublist in regular_list for item in sublist]
    return regular_list

def get_param_index(param, param_list):
    param_list = list(param_list)
    for idx, p in enumerate(param_list):
        if p is param:
            return idx
    print("no corresponding parameter found!")

def get_multiplier(problem):
    if problem.leaf:
        return 1

    assert len(problem.children) > 0
    # stack to store all the nodes of tree
    s1 = []
    # stack to store all the leaf nodes
    s2 = []

    s1.append((problem, 1))
    while len(s1) != 0:
        curr, multiplier = s1.pop(0)

        if len(curr.children) != 0:
            for child in curr.children:
                s1.append((child, multiplier * curr.config.step))
        else:
            s2.append(multiplier)

    assert all(x == s2[0] for x in s2)
    return s2[0]

def log_from_loss_dict(loss_dict):
    outputs = []
    for key, values in loss_dict.items():
        if isinstance(values, dict) or isinstance(values, list):
            for value_idx, value in enumerate(values):
                full_key = key + "_" + str(value_idx)
                if torch.is_tensor(value):
                    value = value.item()
                output = f"{full_key}: {value}"
                outputs.append(output)
        else:
            if torch.is_tensor(values):
                values = values.item()
            output = f"{key}: {values}"
            outputs.append(output)
    return " || ".join(outputs)

def to_vec(tensor_list, alpha=1.0):
    return torch.cat([alpha * t.reshape(-1) for t in tensor_list])

def count_parameters(tensor_list):
    return sum([tensor.numel() for tensor in tensor_list])

def neg_with_none(a):
    if a is None:
        return None
    else:
        return -a

def replace_none_with_zero(tensor_list, reference):

    out = []
    for t, r in zip(tensor_list, reference):
        fixed = t if t is not None else torch.zeros_like(r)
        out.append(fixed)
    return tuple(out)


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """

        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """
        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )

    def train_step(self):
        """
        Running one-step gradient descent for all leaf problems.
        """
        for leaf in self.leaves:
            leaf.step(global_step=self.global_step)

    def run(self):
        """
        Execute multilevel optimization by running gradient descent for leaf problems.
        """
        self.train()
        for it in range(1, self.train_iters + 1):
            self.global_step += 1
            self.train_step()

            if it % self.valid_step == 0 and self.do_validation():
                self.eval()
                validation_stats = self.validation() or {}
                log_loss = log_from_loss_dict(validation_stats)
                self.logger.info(
                    f"[Validation] [Global Step {self.global_step}] " f"{log_loss}"
                )
                self.logger.log(
                    validation_stats, tag="validation", step=self.global_step
                )
                self.train()

                # early stopping
                if self.early_stopping is not None:
                    stop = self.early_stopping(validation_stats)
                    if stop:
                        self.logger.info("Early stopping is executed!")
                        break

    def initialize(self):
        """
        Initialize dependencies (computational graph) between problems.
        """
        # Parse config
        self.parse_config()

        # initialize distributed training
        dist_dict = self.configure_systems()

        # initialize logger
        self.logger = logger(logger_type=self.logger_type)
        if self.is_rank_zero():
            self.logger.info("Initializing Multilevel Optimization...\n")
        start = time.time()

        # parse problem dependency
        self.parse_dependency()

        # set problem attributes
        for problem in self.problems:
            self.set_problem_attr(problem)

        # env initialization
        if self.env is not None:
            self.env.configure_distributed_training(dist_dict)
            self.env.configure_device(self.device)
            self.env.initialize()

        # problem initialization
        for problem in self.problems:
            problem.add_logger(self.logger)
            problem.configure_distributed_training(dist_dict)
            problem.configure_device(self.device)
            problem.configure_roll_back(self._roll_back)
            problem.initialize()
            if self.env is not None:
                problem.add_env(self.env)

        end = time.time()
        if self.is_rank_zero():
            self.logger.info(f"Time spent on initialization: {end-start:.3f} (s)\n")

    def configure_systems(self):
        """
        Configure basic systems set-up like distributed training and device placement.
        """
        # configure distributed training
        if self._strategy in ["distributed", "zero", "fsdp"]:
            dist.init_process_group(backend=self._backend)

            self._world_size = dist.get_world_size()
            assert self._world_size > 1
            self._rank = dist.get_rank()

            device_count = torch.cuda.device_count()
            self._local_rank = self._rank % device_count

        dist_dict = {}
        dist_dict["strategy"] = self._strategy
        dist_dict["backend"] = self._backend
        dist_dict["world_size"] = self._world_size
        dist_dict["rank"] = self._rank
        dist_dict["local_rank"] = self._local_rank

        # configure device for the current rank
        if self._strategy in ["distributed", "zero", "fsdp"]:
            torch.cuda.set_device(self._local_rank)
            self.device = torch.device("cuda", self._local_rank)
        elif self._strategy == "accelerate":
            self.device = self.accelerator.device
        elif self._strategy == "cpu":
            self.device = "cpu"
        elif self._strategy == "gpu":
            self.device = "cuda"
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        return dist_dict

    def train(self):
        """
        Set all problems in multilevel optimization to the train mode.
        """
        for problem in self.problems:
            problem.train()

    def eval(self):
        """
        Set all problems in multilevel optimization to the eval mode.
        """
        for problem in self.problems:
            problem.eval()

    def check_leaf(self, problem):
        """
        Check whether the given ``problem`` is a leaf problem or not.

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: True or False
        :rtype: bool
        """

        for _, value_list in self.dependencies["l2u"].items():
            if problem in set(value_list):
                return False

        return True


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """
        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )

    def train_step(self):
        """
        Running one-step gradient descent for all leaf problems.
        """
        for leaf in self.leaves:
            leaf.step(global_step=self.global_step)

    def run(self):
        """
        Execute multilevel optimization by running gradient descent for leaf problems.
        """
        self.train()
        for it in range(1, self.train_iters + 1):
            self.global_step += 1
            self.train_step()

            if it % self.valid_step == 0 and self.do_validation():
                self.eval()
                validation_stats = self.validation() or {}
                log_loss = log_from_loss_dict(validation_stats)
                self.logger.info(
                    f"[Validation] [Global Step {self.global_step}] " f"{log_loss}"
                )
                self.logger.log(
                    validation_stats, tag="validation", step=self.global_step
                )
                self.train()

                # early stopping
                if self.early_stopping is not None:
                    stop = self.early_stopping(validation_stats)
                    if stop:
                        self.logger.info("Early stopping is executed!")
                        break

    def initialize(self):
        """
        Initialize dependencies (computational graph) between problems.
        """
        # Parse config
        self.parse_config()

        # initialize distributed training
        dist_dict = self.configure_systems()

        # initialize logger
        self.logger = logger(logger_type=self.logger_type)
        if self.is_rank_zero():
            self.logger.info("Initializing Multilevel Optimization...\n")
        start = time.time()

        # parse problem dependency
        self.parse_dependency()

        # set problem attributes
        for problem in self.problems:
            self.set_problem_attr(problem)

        # env initialization
        if self.env is not None:
            self.env.configure_distributed_training(dist_dict)
            self.env.configure_device(self.device)
            self.env.initialize()

        # problem initialization
        for problem in self.problems:
            problem.add_logger(self.logger)
            problem.configure_distributed_training(dist_dict)
            problem.configure_device(self.device)
            problem.configure_roll_back(self._roll_back)
            problem.initialize()
            if self.env is not None:
                problem.add_env(self.env)

        end = time.time()
        if self.is_rank_zero():
            self.logger.info(f"Time spent on initialization: {end-start:.3f} (s)\n")

    def configure_systems(self):
        """
        Configure basic systems set-up like distributed training and device placement.
        """
        # configure distributed training
        if self._strategy in ["distributed", "zero", "fsdp"]:
            dist.init_process_group(backend=self._backend)

            self._world_size = dist.get_world_size()
            assert self._world_size > 1
            self._rank = dist.get_rank()

            device_count = torch.cuda.device_count()
            self._local_rank = self._rank % device_count

        dist_dict = {}
        dist_dict["strategy"] = self._strategy
        dist_dict["backend"] = self._backend
        dist_dict["world_size"] = self._world_size
        dist_dict["rank"] = self._rank
        dist_dict["local_rank"] = self._local_rank

        # configure device for the current rank
        if self._strategy in ["distributed", "zero", "fsdp"]:
            torch.cuda.set_device(self._local_rank)
            self.device = torch.device("cuda", self._local_rank)
        elif self._strategy == "accelerate":
            self.device = self.accelerator.device
        elif self._strategy == "cpu":
            self.device = "cpu"
        elif self._strategy == "gpu":
            self.device = "cuda"
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        return dist_dict

    def train(self):
        """
        Set all problems in multilevel optimization to the train mode.
        """
        for problem in self.problems:
            problem.train()

    def eval(self):
        """
        Set all problems in multilevel optimization to the eval mode.
        """
        for problem in self.problems:
            problem.eval()

    def check_leaf(self, problem):
        """
        Check whether the given ``problem`` is a leaf problem or not.

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: True or False
        :rtype: bool
        """
        for _, value_list in self.dependencies["l2u"].items():
            if problem in set(value_list):
                return False

        return True

    def find_paths(self, src, dst):
        """
        Find all paths from ``src`` to ``dst`` with a modified depth-first search algorithm.

        :param src: The end point of the upper-to-lower edge.
        :type src: Problem
        :param dst: The start point of the upper-to-lower edge.
        :type dst: Problem
        :return: List of all paths from ``src`` to ``dst``.
        """

        results = []
        path = [src]
        self.dfs(src, dst, path, results)
        assert len(results) > 0, f"No path from {src.name} to {dst.name}!"

        for i, _ in enumerate(results):
            results[i].reverse()
            results[i].append(dst)

        return results


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """
        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )

    def train_step(self):
        """
        Running one-step gradient descent for all leaf problems.
        """
        for leaf in self.leaves:
            leaf.step(global_step=self.global_step)

    def run(self):
        """
        Execute multilevel optimization by running gradient descent for leaf problems.
        """
        self.train()
        for it in range(1, self.train_iters + 1):
            self.global_step += 1
            self.train_step()

            if it % self.valid_step == 0 and self.do_validation():
                self.eval()
                validation_stats = self.validation() or {}
                log_loss = log_from_loss_dict(validation_stats)
                self.logger.info(
                    f"[Validation] [Global Step {self.global_step}] " f"{log_loss}"
                )
                self.logger.log(
                    validation_stats, tag="validation", step=self.global_step
                )
                self.train()

                # early stopping
                if self.early_stopping is not None:
                    stop = self.early_stopping(validation_stats)
                    if stop:
                        self.logger.info("Early stopping is executed!")
                        break

    def initialize(self):
        """
        Initialize dependencies (computational graph) between problems.
        """
        # Parse config
        self.parse_config()

        # initialize distributed training
        dist_dict = self.configure_systems()

        # initialize logger
        self.logger = logger(logger_type=self.logger_type)
        if self.is_rank_zero():
            self.logger.info("Initializing Multilevel Optimization...\n")
        start = time.time()

        # parse problem dependency
        self.parse_dependency()

        # set problem attributes
        for problem in self.problems:
            self.set_problem_attr(problem)

        # env initialization
        if self.env is not None:
            self.env.configure_distributed_training(dist_dict)
            self.env.configure_device(self.device)
            self.env.initialize()

        # problem initialization
        for problem in self.problems:
            problem.add_logger(self.logger)
            problem.configure_distributed_training(dist_dict)
            problem.configure_device(self.device)
            problem.configure_roll_back(self._roll_back)
            problem.initialize()
            if self.env is not None:
                problem.add_env(self.env)

        end = time.time()
        if self.is_rank_zero():
            self.logger.info(f"Time spent on initialization: {end-start:.3f} (s)\n")

    def configure_systems(self):
        """
        Configure basic systems set-up like distributed training and device placement.
        """
        # configure distributed training
        if self._strategy in ["distributed", "zero", "fsdp"]:
            dist.init_process_group(backend=self._backend)

            self._world_size = dist.get_world_size()
            assert self._world_size > 1
            self._rank = dist.get_rank()

            device_count = torch.cuda.device_count()
            self._local_rank = self._rank % device_count

        dist_dict = {}
        dist_dict["strategy"] = self._strategy
        dist_dict["backend"] = self._backend
        dist_dict["world_size"] = self._world_size
        dist_dict["rank"] = self._rank
        dist_dict["local_rank"] = self._local_rank

        # configure device for the current rank
        if self._strategy in ["distributed", "zero", "fsdp"]:
            torch.cuda.set_device(self._local_rank)
            self.device = torch.device("cuda", self._local_rank)
        elif self._strategy == "accelerate":
            self.device = self.accelerator.device
        elif self._strategy == "cpu":
            self.device = "cpu"
        elif self._strategy == "gpu":
            self.device = "cuda"
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        return dist_dict

    def train(self):
        """
        Set all problems in multilevel optimization to the train mode.
        """
        for problem in self.problems:
            problem.train()

    def eval(self):
        """
        Set all problems in multilevel optimization to the eval mode.
        """
        for problem in self.problems:
            problem.eval()

    def check_leaf(self, problem):
        """
        Check whether the given ``problem`` is a leaf problem or not.

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: True or False
        :rtype: bool
        """
        for _, value_list in self.dependencies["l2u"].items():
            if problem in set(value_list):
                return False

        return True

    def find_paths(self, src, dst):
        """
        Find all paths from ``src`` to ``dst`` with a modified depth-first search algorithm.

        :param src: The end point of the upper-to-lower edge.
        :type src: Problem
        :param dst: The start point of the upper-to-lower edge.
        :type dst: Problem
        :return: List of all paths from ``src`` to ``dst``.
        """
        results = []
        path = [src]
        self.dfs(src, dst, path, results)
        assert len(results) > 0, f"No path from {src.name} to {dst.name}!"

        for i, _ in enumerate(results):
            results[i].reverse()
            results[i].append(dst)

        return results

    def dfs(self, src, dst, path, results):

        if src is dst:
            assert len(path) > 1
            result = [node for node in path]
            results.append(result)
        elif src not in self.dependencies["l2u"]:
            return
        else:
            for adj in self.dependencies["l2u"][src]:
                path.append(adj)
                self.dfs(adj, dst, path, results)
                path.pop()


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """
        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )

    def train_step(self):
        """
        Running one-step gradient descent for all leaf problems.
        """
        for leaf in self.leaves:
            leaf.step(global_step=self.global_step)

    def run(self):
        """
        Execute multilevel optimization by running gradient descent for leaf problems.
        """
        self.train()
        for it in range(1, self.train_iters + 1):
            self.global_step += 1
            self.train_step()

            if it % self.valid_step == 0 and self.do_validation():
                self.eval()
                validation_stats = self.validation() or {}
                log_loss = log_from_loss_dict(validation_stats)
                self.logger.info(
                    f"[Validation] [Global Step {self.global_step}] " f"{log_loss}"
                )
                self.logger.log(
                    validation_stats, tag="validation", step=self.global_step
                )
                self.train()

                # early stopping
                if self.early_stopping is not None:
                    stop = self.early_stopping(validation_stats)
                    if stop:
                        self.logger.info("Early stopping is executed!")
                        break

    def initialize(self):
        """
        Initialize dependencies (computational graph) between problems.
        """
        # Parse config
        self.parse_config()

        # initialize distributed training
        dist_dict = self.configure_systems()

        # initialize logger
        self.logger = logger(logger_type=self.logger_type)
        if self.is_rank_zero():
            self.logger.info("Initializing Multilevel Optimization...\n")
        start = time.time()

        # parse problem dependency
        self.parse_dependency()

        # set problem attributes
        for problem in self.problems:
            self.set_problem_attr(problem)

        # env initialization
        if self.env is not None:
            self.env.configure_distributed_training(dist_dict)
            self.env.configure_device(self.device)
            self.env.initialize()

        # problem initialization
        for problem in self.problems:
            problem.add_logger(self.logger)
            problem.configure_distributed_training(dist_dict)
            problem.configure_device(self.device)
            problem.configure_roll_back(self._roll_back)
            problem.initialize()
            if self.env is not None:
                problem.add_env(self.env)

        end = time.time()
        if self.is_rank_zero():
            self.logger.info(f"Time spent on initialization: {end-start:.3f} (s)\n")

    def configure_systems(self):
        """
        Configure basic systems set-up like distributed training and device placement.
        """
        # configure distributed training
        if self._strategy in ["distributed", "zero", "fsdp"]:
            dist.init_process_group(backend=self._backend)

            self._world_size = dist.get_world_size()
            assert self._world_size > 1
            self._rank = dist.get_rank()

            device_count = torch.cuda.device_count()
            self._local_rank = self._rank % device_count

        dist_dict = {}
        dist_dict["strategy"] = self._strategy
        dist_dict["backend"] = self._backend
        dist_dict["world_size"] = self._world_size
        dist_dict["rank"] = self._rank
        dist_dict["local_rank"] = self._local_rank

        # configure device for the current rank
        if self._strategy in ["distributed", "zero", "fsdp"]:
            torch.cuda.set_device(self._local_rank)
            self.device = torch.device("cuda", self._local_rank)
        elif self._strategy == "accelerate":
            self.device = self.accelerator.device
        elif self._strategy == "cpu":
            self.device = "cpu"
        elif self._strategy == "gpu":
            self.device = "cuda"
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        return dist_dict

    def train(self):
        """
        Set all problems in multilevel optimization to the train mode.
        """
        for problem in self.problems:
            problem.train()

    def eval(self):
        """
        Set all problems in multilevel optimization to the eval mode.
        """
        for problem in self.problems:
            problem.eval()

    def check_leaf(self, problem):
        """
        Check whether the given ``problem`` is a leaf problem or not.

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: True or False
        :rtype: bool
        """
        for _, value_list in self.dependencies["l2u"].items():
            if problem in set(value_list):
                return False

        return True

    def find_paths(self, src, dst):
        """
        Find all paths from ``src`` to ``dst`` with a modified depth-first search algorithm.

        :param src: The end point of the upper-to-lower edge.
        :type src: Problem
        :param dst: The start point of the upper-to-lower edge.
        :type dst: Problem
        :return: List of all paths from ``src`` to ``dst``.
        """
        results = []
        path = [src]
        self.dfs(src, dst, path, results)
        assert len(results) > 0, f"No path from {src.name} to {dst.name}!"

        for i, _ in enumerate(results):
            results[i].reverse()
            results[i].append(dst)

        return results

    def dfs(self, src, dst, path, results):
        if src is dst:
            assert len(path) > 1
            result = [node for node in path]
            results.append(result)
        elif src not in self.dependencies["l2u"]:
            return
        else:
            for adj in self.dependencies["l2u"][src]:
                path.append(adj)
                self.dfs(adj, dst, path, results)
                path.pop()

    def parse_dependency(self):
        """
        Parse user-provided ``u2l`` and ``l2u`` dependencies to figure out 1) topological order for
        multilevel optimization execution, and 2) backpropagation path(s) for each problem. A
        modified depth-first search algorithm is used.
        """
        # Parse upper-to-lower dependency

        for key, value_list in self.dependencies["u2l"].items():
            for value in value_list:
                # find all paths from low to high for backpropagation
                paths = self.find_paths(src=value, dst=key)
                key.add_paths(paths)

        # Parse lower-to-upper dependency
        for key, value_list in self.dependencies["l2u"].items():
            for value in value_list:
                # add value problem to parents of key problem for backpropgation
                key.add_parent(value)
                value.add_child(key)

        # Parse problems
        for problem in self.problems:
            if self.check_leaf(problem):
                problem.leaf = True
                self.leaves.append(problem)


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """
        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )

    def train_step(self):
        """
        Running one-step gradient descent for all leaf problems.
        """
        for leaf in self.leaves:
            leaf.step(global_step=self.global_step)

    def run(self):
        """
        Execute multilevel optimization by running gradient descent for leaf problems.
        """
        self.train()
        for it in range(1, self.train_iters + 1):
            self.global_step += 1
            self.train_step()

            if it % self.valid_step == 0 and self.do_validation():
                self.eval()
                validation_stats = self.validation() or {}
                log_loss = log_from_loss_dict(validation_stats)
                self.logger.info(
                    f"[Validation] [Global Step {self.global_step}] " f"{log_loss}"
                )
                self.logger.log(
                    validation_stats, tag="validation", step=self.global_step
                )
                self.train()

                # early stopping
                if self.early_stopping is not None:
                    stop = self.early_stopping(validation_stats)
                    if stop:
                        self.logger.info("Early stopping is executed!")
                        break

    def initialize(self):
        """
        Initialize dependencies (computational graph) between problems.
        """
        # Parse config
        self.parse_config()

        # initialize distributed training
        dist_dict = self.configure_systems()

        # initialize logger
        self.logger = logger(logger_type=self.logger_type)
        if self.is_rank_zero():
            self.logger.info("Initializing Multilevel Optimization...\n")
        start = time.time()

        # parse problem dependency
        self.parse_dependency()

        # set problem attributes
        for problem in self.problems:
            self.set_problem_attr(problem)

        # env initialization
        if self.env is not None:
            self.env.configure_distributed_training(dist_dict)
            self.env.configure_device(self.device)
            self.env.initialize()

        # problem initialization
        for problem in self.problems:
            problem.add_logger(self.logger)
            problem.configure_distributed_training(dist_dict)
            problem.configure_device(self.device)
            problem.configure_roll_back(self._roll_back)
            problem.initialize()
            if self.env is not None:
                problem.add_env(self.env)

        end = time.time()
        if self.is_rank_zero():
            self.logger.info(f"Time spent on initialization: {end-start:.3f} (s)\n")

    def configure_systems(self):
        """
        Configure basic systems set-up like distributed training and device placement.
        """
        # configure distributed training
        if self._strategy in ["distributed", "zero", "fsdp"]:
            dist.init_process_group(backend=self._backend)

            self._world_size = dist.get_world_size()
            assert self._world_size > 1
            self._rank = dist.get_rank()

            device_count = torch.cuda.device_count()
            self._local_rank = self._rank % device_count

        dist_dict = {}
        dist_dict["strategy"] = self._strategy
        dist_dict["backend"] = self._backend
        dist_dict["world_size"] = self._world_size
        dist_dict["rank"] = self._rank
        dist_dict["local_rank"] = self._local_rank

        # configure device for the current rank
        if self._strategy in ["distributed", "zero", "fsdp"]:
            torch.cuda.set_device(self._local_rank)
            self.device = torch.device("cuda", self._local_rank)
        elif self._strategy == "accelerate":
            self.device = self.accelerator.device
        elif self._strategy == "cpu":
            self.device = "cpu"
        elif self._strategy == "gpu":
            self.device = "cuda"
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        return dist_dict

    def train(self):
        """
        Set all problems in multilevel optimization to the train mode.
        """
        for problem in self.problems:
            problem.train()

    def eval(self):
        """
        Set all problems in multilevel optimization to the eval mode.
        """
        for problem in self.problems:
            problem.eval()

    def check_leaf(self, problem):
        """
        Check whether the given ``problem`` is a leaf problem or not.

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: True or False
        :rtype: bool
        """
        for _, value_list in self.dependencies["l2u"].items():
            if problem in set(value_list):
                return False

        return True

    def find_paths(self, src, dst):
        """
        Find all paths from ``src`` to ``dst`` with a modified depth-first search algorithm.

        :param src: The end point of the upper-to-lower edge.
        :type src: Problem
        :param dst: The start point of the upper-to-lower edge.
        :type dst: Problem
        :return: List of all paths from ``src`` to ``dst``.
        """
        results = []
        path = [src]
        self.dfs(src, dst, path, results)
        assert len(results) > 0, f"No path from {src.name} to {dst.name}!"

        for i, _ in enumerate(results):
            results[i].reverse()
            results[i].append(dst)

        return results

    def dfs(self, src, dst, path, results):
        if src is dst:
            assert len(path) > 1
            result = [node for node in path]
            results.append(result)
        elif src not in self.dependencies["l2u"]:
            return
        else:
            for adj in self.dependencies["l2u"][src]:
                path.append(adj)
                self.dfs(adj, dst, path, results)
                path.pop()

    def parse_dependency(self):
        """
        Parse user-provided ``u2l`` and ``l2u`` dependencies to figure out 1) topological order for
        multilevel optimization execution, and 2) backpropagation path(s) for each problem. A
        modified depth-first search algorithm is used.
        """
        # Parse upper-to-lower dependency
        for key, value_list in self.dependencies["u2l"].items():
            for value in value_list:
                # find all paths from low to high for backpropagation
                paths = self.find_paths(src=value, dst=key)
                key.add_paths(paths)

        # Parse lower-to-upper dependency
        for key, value_list in self.dependencies["l2u"].items():
            for value in value_list:
                # add value problem to parents of key problem for backpropgation
                key.add_parent(value)
                value.add_child(key)

        # Parse problems
        for problem in self.problems:
            if self.check_leaf(problem):
                problem.leaf = True
                self.leaves.append(problem)

    def set_dependency(self, dependencies):
        self.dependencies = dependencies
        self.leaves = []

        # clear existing dependencies
        for problem in self.problems:
            problem.leaf = False
            problem.clear_dependencies()

        self.parse_dependency()

    def set_problem_attr(self, problem):
        """
        Set class attribute for the given ``problem`` based on their names

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: ``problem`` name
        :rtype: str
        """

        name = problem.name

        # set attribute for Engine
        assert not hasattr(self, name), f"Problem already has a problelm named {name}!"
        setattr(self, name, problem)

        # set attribute for Problems
        for prob in self.problems:
            if prob != problem:
                assert not hasattr(problem, name)
                setattr(prob, name, problem)

        # set attribute for Env
        if self.env is not None:
            setattr(self.env, name, problem)

        return name


Prompt:
# Copyright Sang Keun Choe
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import torch
import torch.distributed as dist

from betty.configs import EngineConfig
from betty.logging import logger
from betty.misc.early_stopping import EarlyStopping
from betty.utils import log_from_loss_dict

class Engine:
    """
    ``Engine`` handles a dataflow graph based on the user-provided hierarchical problem
    dependencies. It also provides a primitive for executing multilevel optimization.
    """

    def __init__(self, problems, config=None, dependencies=None, env=None):
        # config
        self.config = config if config is not None else EngineConfig()

        # step counters
        self.train_iters = 0
        self.valid_step = 0
        self.global_step = 0

        # logger
        self.logger_type = None
        self.logger = None

        # problem
        self.problems = problems
        self.leaves = []

        # dependencies
        self.dependencies = dependencies

        # env
        self.env = env

        # distributed
        self._strategy = None
        self._backend = None
        self._world_size = 0
        self._rank = 0
        self._local_rank = 0

        # early stopping
        self.early_stopping = None

        # roll back
        self._roll_back = False

        # device
        self.device = None

        # initialize
        self.initialize()

    def parse_config(self):
        """
        Parse EngineConfig.
        """
        self.train_iters = self.config.train_iters
        self.valid_step = self.config.valid_step

        self.logger_type = self.config.logger_type

        self._roll_back = self.config.roll_back

        self._strategy = self.config.strategy
        self._backend = self.config.backend

        if self.config.early_stopping:
            self.early_stopping = EarlyStopping(
                metric=self.config.early_stopping_metric,
                mode=self.config.early_stopping_mode,
                tolerance=self.config.early_stopping_tolerance,
            )

    def train_step(self):
        """
        Running one-step gradient descent for all leaf problems.
        """
        for leaf in self.leaves:
            leaf.step(global_step=self.global_step)

    def run(self):
        """
        Execute multilevel optimization by running gradient descent for leaf problems.
        """
        self.train()
        for it in range(1, self.train_iters + 1):
            self.global_step += 1
            self.train_step()

            if it % self.valid_step == 0 and self.do_validation():
                self.eval()
                validation_stats = self.validation() or {}
                log_loss = log_from_loss_dict(validation_stats)
                self.logger.info(
                    f"[Validation] [Global Step {self.global_step}] " f"{log_loss}"
                )
                self.logger.log(
                    validation_stats, tag="validation", step=self.global_step
                )
                self.train()

                # early stopping
                if self.early_stopping is not None:
                    stop = self.early_stopping(validation_stats)
                    if stop:
                        self.logger.info("Early stopping is executed!")
                        break

    def initialize(self):
        """
        Initialize dependencies (computational graph) between problems.
        """
        # Parse config
        self.parse_config()

        # initialize distributed training
        dist_dict = self.configure_systems()

        # initialize logger
        self.logger = logger(logger_type=self.logger_type)
        if self.is_rank_zero():
            self.logger.info("Initializing Multilevel Optimization...\n")
        start = time.time()

        # parse problem dependency
        self.parse_dependency()

        # set problem attributes
        for problem in self.problems:
            self.set_problem_attr(problem)

        # env initialization
        if self.env is not None:
            self.env.configure_distributed_training(dist_dict)
            self.env.configure_device(self.device)
            self.env.initialize()

        # problem initialization
        for problem in self.problems:
            problem.add_logger(self.logger)
            problem.configure_distributed_training(dist_dict)
            problem.configure_device(self.device)
            problem.configure_roll_back(self._roll_back)
            problem.initialize()
            if self.env is not None:
                problem.add_env(self.env)

        end = time.time()
        if self.is_rank_zero():
            self.logger.info(f"Time spent on initialization: {end-start:.3f} (s)\n")

    def configure_systems(self):
        """
        Configure basic systems set-up like distributed training and device placement.
        """
        # configure distributed training
        if self._strategy in ["distributed", "zero", "fsdp"]:
            dist.init_process_group(backend=self._backend)

            self._world_size = dist.get_world_size()
            assert self._world_size > 1
            self._rank = dist.get_rank()

            device_count = torch.cuda.device_count()
            self._local_rank = self._rank % device_count

        dist_dict = {}
        dist_dict["strategy"] = self._strategy
        dist_dict["backend"] = self._backend
        dist_dict["world_size"] = self._world_size
        dist_dict["rank"] = self._rank
        dist_dict["local_rank"] = self._local_rank

        # configure device for the current rank
        if self._strategy in ["distributed", "zero", "fsdp"]:
            torch.cuda.set_device(self._local_rank)
            self.device = torch.device("cuda", self._local_rank)
        elif self._strategy == "accelerate":
            self.device = self.accelerator.device
        elif self._strategy == "cpu":
            self.device = "cpu"
        elif self._strategy == "gpu":
            self.device = "cuda"
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        return dist_dict

    def train(self):
        """
        Set all problems in multilevel optimization to the train mode.
        """
        for problem in self.problems:
            problem.train()

    def eval(self):
        """
        Set all problems in multilevel optimization to the eval mode.
        """
        for problem in self.problems:
            problem.eval()

    def check_leaf(self, problem):
        """
        Check whether the given ``problem`` is a leaf problem or not.

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: True or False
        :rtype: bool
        """
        for _, value_list in self.dependencies["l2u"].items():
            if problem in set(value_list):
                return False

        return True

    def find_paths(self, src, dst):
        """
        Find all paths from ``src`` to ``dst`` with a modified depth-first search algorithm.

        :param src: The end point of the upper-to-lower edge.
        :type src: Problem
        :param dst: The start point of the upper-to-lower edge.
        :type dst: Problem
        :return: List of all paths from ``src`` to ``dst``.
        """
        results = []
        path = [src]
        self.dfs(src, dst, path, results)
        assert len(results) > 0, f"No path from {src.name} to {dst.name}!"

        for i, _ in enumerate(results):
            results[i].reverse()
            results[i].append(dst)

        return results

    def dfs(self, src, dst, path, results):
        if src is dst:
            assert len(path) > 1
            result = [node for node in path]
            results.append(result)
        elif src not in self.dependencies["l2u"]:
            return
        else:
            for adj in self.dependencies["l2u"][src]:
                path.append(adj)
                self.dfs(adj, dst, path, results)
                path.pop()

    def parse_dependency(self):
        """
        Parse user-provided ``u2l`` and ``l2u`` dependencies to figure out 1) topological order for
        multilevel optimization execution, and 2) backpropagation path(s) for each problem. A
        modified depth-first search algorithm is used.
        """
        # Parse upper-to-lower dependency
        for key, value_list in self.dependencies["u2l"].items():
            for value in value_list:
                # find all paths from low to high for backpropagation
                paths = self.find_paths(src=value, dst=key)
                key.add_paths(paths)

        # Parse lower-to-upper dependency
        for key, value_list in self.dependencies["l2u"].items():
            for value in value_list:
                # add value problem to parents of key problem for backpropgation
                key.add_parent(value)
                value.add_child(key)

        # Parse problems
        for problem in self.problems:
            if self.check_leaf(problem):
                problem.leaf = True
                self.leaves.append(problem)

    def set_dependency(self, dependencies):
        self.dependencies = dependencies
        self.leaves = []

        # clear existing dependencies
        for problem in self.problems:
            problem.leaf = False
            problem.clear_dependencies()

        self.parse_dependency()

    def set_problem_attr(self, problem):
        """
        Set class attribute for the given ``problem`` based on their names

        :param problem: Problem in multilevel optimization
        :type problem: Problem
        :return: ``problem`` name
        :rtype: str
        """
        name = problem.name

        # set attribute for Engine
        assert not hasattr(self, name), f"Problem already has a problelm named {name}!"
        setattr(self, name, problem)

        # set attribute for Problems
        for prob in self.problems:
            if prob != problem:
                assert not hasattr(problem, name)
                setattr(prob, name, problem)

        # set attribute for Env
        if self.env is not None:
            setattr(self.env, name, problem)

        return name

    def do_validation(self):
        """
        Check whether to run validation.
        """

        if self.is_implemented("validation") and self.is_rank_zero():
            return True
        return False


Prompt:
from torch.distributed.optim import ZeroRedundancyOptimizer

def patch_optimizer(optimizer, params, is_zero):

    defaults = optimizer.defaults
    new_optimizer = None
    if is_zero:
        new_optimizer = ZeroRedundancyOptimizer(
            params=params,
            optimizer_class=optimizer.__class__,
            parameters_as_bucket_view=True,
            **defaults,
        )
    else:
        new_optimizer = optimizer.__class__(params, **defaults)

    return new_optimizer


Prompt:
import warnings

import torch

from betty.utils import neg_with_none

def neumann(vector, curr, prev, sync):
    """
    Approximate the matrix-vector multiplication with the best response Jacobian by the
    Neumann Series as proposed in
    `Optimizing Millions of Hyperparameters by Implicit Differentiation
    <https://arxiv.org/abs/1911.02590>`_ based on implicit function theorem (IFT). Users may
    specify learning rate (``neumann_alpha``) and unrolling steps (``neumann_iterations``) in
    ``Config``.

    :param vector:
        Vector with which matrix-vector multiplication with best-response Jacobian (matrix) would
        be performed.
    :type vector: Sequence of Tensor
    :param curr: A current level problem
    :type curr: Problem
    :param prev: A directly lower-level problem to the current problem
    :type prev: Problem
    :return: (Intermediate) gradient
    :rtype: Sequence of Tensor
    """
    # ! Mabye replace with child.loss by adding self.loss attribute to save computation

    assert len(curr.paths) == 0, "neumann method is not supported for higher order MLO!"
    config = curr.config
    in_loss = curr.training_step_exec(curr.cur_batch)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        in_grad = torch.autograd.grad(
            in_loss, curr.trainable_parameters(), create_graph=True
        )
    v2 = approx_inverse_hvp(
        vector,
        in_grad,
        curr.trainable_parameters(),
        iterations=config.neumann_iterations,
        alpha=config.neumann_alpha,
    )
    if sync:
        v2 = [neg_with_none(x) for x in v2]
        torch.autograd.backward(
            in_grad, inputs=prev.trainable_parameters(), grad_tensors=v2
        )
        implicit_grad = None
    else:
        implicit_grad = torch.autograd.grad(
            in_grad, prev.trainable_parameters(), grad_outputs=v2
        )
        implicit_grad = [neg_with_none(ig) for ig in implicit_grad]

    return implicit_grad


Prompt:
import warnings

import torch

from betty.utils import neg_with_none

def neumann(vector, curr, prev, sync):
    """
    Approximate the matrix-vector multiplication with the best response Jacobian by the
    Neumann Series as proposed in
    `Optimizing Millions of Hyperparameters by Implicit Differentiation
    <https://arxiv.org/abs/1911.02590>`_ based on implicit function theorem (IFT). Users may
    specify learning rate (``neumann_alpha``) and unrolling steps (``neumann_iterations``) in
    ``Config``.

    :param vector:
        Vector with which matrix-vector multiplication with best-response Jacobian (matrix) would
        be performed.
    :type vector: Sequence of Tensor
    :param curr: A current level problem
    :type curr: Problem
    :param prev: A directly lower-level problem to the current problem
    :type prev: Problem
    :return: (Intermediate) gradient
    :rtype: Sequence of Tensor
    """
    # ! Mabye replace with child.loss by adding self.loss attribute to save computation
    assert len(curr.paths) == 0, "neumann method is not supported for higher order MLO!"
    config = curr.config
    in_loss = curr.training_step_exec(curr.cur_batch)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        in_grad = torch.autograd.grad(
            in_loss, curr.trainable_parameters(), create_graph=True
        )
    v2 = approx_inverse_hvp(
        vector,
        in_grad,
        curr.trainable_parameters(),
        iterations=config.neumann_iterations,
        alpha=config.neumann_alpha,
    )
    if sync:
        v2 = [neg_with_none(x) for x in v2]
        torch.autograd.backward(
            in_grad, inputs=prev.trainable_parameters(), grad_tensors=v2
        )
        implicit_grad = None
    else:
        implicit_grad = torch.autograd.grad(
            in_grad, prev.trainable_parameters(), grad_outputs=v2
        )
        implicit_grad = [neg_with_none(ig) for ig in implicit_grad]

    return implicit_grad

def approx_inverse_hvp(v, f, params, iterations=3, alpha=1.0):

    p = v
    for _ in range(iterations):
        hvp = torch.autograd.grad(f, params, grad_outputs=v, retain_graph=True)
        v = [v_i - alpha * hvp_i for v_i, hvp_i in zip(v, hvp)]
        p = [v_i + p_i for v_i, p_i in zip(v, p)]

    return [alpha * p_i for p_i in p]


Prompt:
import math
import torch

def get_optimzer_type(optimizer):

    cls_name = type(optimizer).__name__.lower()
    if "adam" in cls_name:
        return "adam"
    return "sgd"


Prompt:
import torch

from betty.utils import replace_none_with_zero

from .darts import darts
from .cg import cg
from .neumann import neumann
from .reinforce import reinforce

jvp_fn_mapping = {
    "darts": darts,
    "neumann": neumann,
    "cg": cg,
    "reinforce": reinforce,
}

def get_grads(loss, path, retain_graph, do_sync):

    jvp = torch.autograd.grad(
        loss,
        path[1].trainable_parameters(),
        retain_graph=retain_graph,
        allow_unused=True,
    )
    jvp = replace_none_with_zero(jvp, path[1].trainable_parameters())
    for i in range(1, len(path) - 1):
        jvp_fn_type = path[i].config.type
        assert jvp_fn_type in jvp_fn_mapping
        jvp_fn = jvp_fn_mapping[jvp_fn_type]
        sync = bool(do_sync and i == len(path) - 2)
        jvp = jvp_fn(jvp, path[i], path[i + 1], sync)

    return jvp


Prompt:
import numpy as np

import torch
import torch.nn.functional as F

from betty.engine import Engine
from betty.configs import Config, EngineConfig
from betty.problems import ImplicitProblem

# hyperparameters
DATA_NUM = 1000
DATA_DIM = 20

# data preparation
w_gt = np.random.randn(DATA_DIM)
x = np.random.randn(DATA_NUM, DATA_DIM)
y = x @ w_gt + 0.1 * np.random.randn(DATA_NUM)
y = (y > 0).astype(float)

idx = DATA_NUM // 2
x_train, x_val, y_train, y_val = x[:idx, :], x[idx:, :], y[:idx], y[idx:]
x_train, y_train = (
    torch.from_numpy(x_train).float(),
    torch.from_numpy(y_train).float(),
)
x_val, y_val = (
    torch.from_numpy(x_val).float(),
    torch.from_numpy(y_val).float(),
)

def make_data_loader(xs, ys):

    datasets = [(xs, ys)]

    return datasets

    datasets = [(xs, ys)]

    return datasets


Prompt:
import numpy as np

import torch
import torch.nn.functional as F

from betty.engine import Engine
from betty.configs import Config, EngineConfig
from betty.problems import ImplicitProblem

# hyperparameters
DATA_NUM = 1000
DATA_DIM = 20

# data preparation
w_gt = np.random.randn(DATA_DIM)
x = np.random.randn(DATA_NUM, DATA_DIM)
y = x @ w_gt + 0.1 * np.random.randn(DATA_NUM)
y = (y > 0).astype(float)

idx = DATA_NUM // 2
x_train, x_val, y_train, y_val = x[:idx, :], x[idx:, :], y[:idx], y[idx:]
x_train, y_train = (
    torch.from_numpy(x_train).float(),
    torch.from_numpy(y_train).float(),
)
x_val, y_val = (
    torch.from_numpy(x_val).float(),
    torch.from_numpy(y_val).float(),
)

def make_data_loader(xs, ys):
    datasets = [(xs, ys)]

    return datasets

class ChildNet(torch.nn.Module):
    def __init__(self) -> None:

        super().__init__()

        self.w = torch.nn.Parameter(torch.zeros(DATA_DIM))


Prompt:
import numpy as np

import torch
import torch.nn.functional as F

from betty.engine import Engine
from betty.configs import Config, EngineConfig
from betty.problems import ImplicitProblem

# hyperparameters
DATA_NUM = 1000
DATA_DIM = 20

# data preparation
w_gt = np.random.randn(DATA_DIM)
x = np.random.randn(DATA_NUM, DATA_DIM)
y = x @ w_gt + 0.1 * np.random.randn(DATA_NUM)
y = (y > 0).astype(float)

idx = DATA_NUM // 2
x_train, x_val, y_train, y_val = x[:idx, :], x[idx:, :], y[:idx], y[idx:]
x_train, y_train = (
    torch.from_numpy(x_train).float(),
    torch.from_numpy(y_train).float(),
)
x_val, y_val = (
    torch.from_numpy(x_val).float(),
    torch.from_numpy(y_val).float(),
)

def make_data_loader(xs, ys):
    datasets = [(xs, ys)]

    return datasets

class ChildNet(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.w = torch.nn.Parameter(torch.zeros(DATA_DIM))

    def forward(self, inputs):
        outs = inputs @ self.w
        return outs, self.w

class ParentNet(torch.nn.Module):
    def __init__(self) -> None:

        super().__init__()

        self.w = torch.nn.Parameter(torch.ones(DATA_DIM))


Prompt:
import numpy as np

import torch
import torch.nn.functional as F

from betty.engine import Engine
from betty.configs import Config, EngineConfig
from betty.problems import ImplicitProblem

# hyperparameters
DATA_NUM = 1000
DATA_DIM = 20

# data preparation
w_gt = np.random.randn(DATA_DIM)
x = np.random.randn(DATA_NUM, DATA_DIM)
y = x @ w_gt + 0.1 * np.random.randn(DATA_NUM)
y = (y > 0).astype(float)

idx = DATA_NUM // 2
x_train, x_val, y_train, y_val = x[:idx, :], x[idx:, :], y[:idx], y[idx:]
x_train, y_train = (
    torch.from_numpy(x_train).float(),
    torch.from_numpy(y_train).float(),
)
x_val, y_val = (
    torch.from_numpy(x_val).float(),
    torch.from_numpy(y_val).float(),
)

def make_data_loader(xs, ys):
    datasets = [(xs, ys)]

    return datasets

class ChildNet(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.w = torch.nn.Parameter(torch.zeros(DATA_DIM))

    def forward(self, inputs):
        outs = inputs @ self.w
        return outs, self.w

class ParentNet(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.w = torch.nn.Parameter(torch.ones(DATA_DIM))

    def forward(self):
        return self.w

class Parent(ImplicitProblem):
    def training_step(self, batch):

        inputs, targets = batch
        outs = self.inner(inputs)[0]
        loss = F.binary_cross_entropy_with_logits(outs, targets)
        return loss


Prompt:
import numpy as np

import torch
import torch.nn.functional as F

from betty.engine import Engine
from betty.configs import Config, EngineConfig
from betty.problems import ImplicitProblem

# hyperparameters
DATA_NUM = 1000
DATA_DIM = 20

# data preparation
w_gt = np.random.randn(DATA_DIM)
x = np.random.randn(DATA_NUM, DATA_DIM)
y = x @ w_gt + 0.1 * np.random.randn(DATA_NUM)
y = (y > 0).astype(float)

idx = DATA_NUM // 2
x_train, x_val, y_train, y_val = x[:idx, :], x[idx:, :], y[:idx], y[idx:]
x_train, y_train = (
    torch.from_numpy(x_train).float(),
    torch.from_numpy(y_train).float(),
)
x_val, y_val = (
    torch.from_numpy(x_val).float(),
    torch.from_numpy(y_val).float(),
)

def make_data_loader(xs, ys):
    datasets = [(xs, ys)]

    return datasets

class ChildNet(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.w = torch.nn.Parameter(torch.zeros(DATA_DIM))

    def forward(self, inputs):
        outs = inputs @ self.w
        return outs, self.w

class ParentNet(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.w = torch.nn.Parameter(torch.ones(DATA_DIM))

    def forward(self):
        return self.w

class Parent(ImplicitProblem):
    def training_step(self, batch):
        inputs, targets = batch
        outs = self.inner(inputs)[0]
        loss = F.binary_cross_entropy_with_logits(outs, targets)
        return loss

    def configure_train_data_loader(self):
        return make_data_loader(x_val, y_val)

    def configure_module(self):
        return ParentNet()

    def configure_optimizer(self):
        return torch.optim.SGD(self.module.parameters(), lr=1, momentum=0.9)

    def param_callback(self):
        for p in self.trainable_parameters():
            p.data.clamp_(min=1e-8)

class Child(ImplicitProblem):
    def training_step(self, batch):

        inputs, targets = batch
        outs, params = self.module(inputs)
        loss = (
            F.binary_cross_entropy_with_logits(outs, targets)
            + 0.5
            * (
                params.unsqueeze(0) @ torch.diag(self.outer()) @ params.unsqueeze(1)
            ).sum()
        )
        return loss


Prompt:
import sys
import abc
import logging

_logger = None

def get_logger():
    """
    Get global logger.
    """

    global _logger
    if _logger:
        return _logger
    logger = logging.getLogger("betty")
    log_format = logging.Formatter(
        "[%(asctime)s] [%(levelname)s] %(message)s", "%Y-%m-%d %H:%M:%S"
    )

    logger.propagate = False
    logger.setLevel(logging.INFO)
    ch = logging.StreamHandler(stream=sys.stdout)
    ch.setLevel(logging.INFO)
    ch.setFormatter(log_format)
    logger.addHandler(ch)

    _logger = logger
    return _logger


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Plotting functions pre and post model fitting."""

import functools
import logging

# Using these types from typing instead of their generic types in the type hints
# in order to be compatible with Python 3.7 and 3.8.
from typing import Any, List, Optional, Sequence, Tuple

import arviz
import jax
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import pandas as pd
import seaborn as sns
from sklearn import metrics

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models
from lightweight_mmm import preprocessing
from lightweight_mmm import utils

plt.style.use("default")

_PALETTE = sns.color_palette(n_colors=100)

@functools.partial(jax.jit, static_argnames=("media_mix_model"))
def _make_single_prediction(media_mix_model: lightweight_mmm.LightweightMMM,
                            mock_media: jnp.ndarray,
                            extra_features: Optional[jnp.ndarray],
                            seed: Optional[int]
                            ) -> jnp.ndarray:
  """Makes a prediction of a single row.

  Serves as a helper function for making predictions individually for each media
  channel and one row at a time. It is meant to be used vmaped otherwise it can
  be slow as it's meant to be used for plotting curve responses only. Use
  lightweight_mmm.LightweightMMM for regular predict functionality.

  Args:
    media_mix_model: Media mix model to use for getting the predictions.
    mock_media: Mock media for this iteration of predictions.
    extra_features: Extra features to use for predictions.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    A point estimate for the given data.
  """

  return media_mix_model.predict(
      media=jnp.expand_dims(mock_media, axis=0),
      extra_features=extra_features,
      seed=seed).mean(axis=0)