Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params_no_mlps = self._zero_mlps(params)

    out_zero_activation = forward_zero.apply(params, next(rng), emb, mask)
    out_no_mlps = forward.apply(params_no_mlps, next(rng), emb, mask)

    self._check_layer_naming(params)
    np.testing.assert_allclose(out_zero_activation, out_no_mlps)
    self.assertFalse(np.allclose(out_zero_activation, 0))

class CompiledTransformerModelTest(parameterized.TestCase):

  def _get_one_hot_embed_unembed(self, vocab_size, max_seq_len):
    # Embeds tokens as one-hot into the first `vocab_size` dimensions

    token_embed = hk.Embed(
        embedding_matrix=jnp.block(
            [jnp.eye(vocab_size),
             jnp.zeros((vocab_size, max_seq_len))]))

    # Embeds positions as one-hot into the last `max_seq_len` dimensions
    position_embed = hk.Embed(
        embedding_matrix=jnp.block(
            [jnp.zeros((max_seq_len, vocab_size)),
             jnp.eye(max_seq_len)]))

    class Unembed(hk.Module):

      def __call__(self, embeddings):
        return jnp.argmax(embeddings[:, :, :vocab_size], axis=-1)

    return token_embed, position_embed, Unembed()


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params_no_mlps = self._zero_mlps(params)

    out_zero_activation = forward_zero.apply(params, next(rng), emb, mask)
    out_no_mlps = forward.apply(params_no_mlps, next(rng), emb, mask)

    self._check_layer_naming(params)
    np.testing.assert_allclose(out_zero_activation, out_no_mlps)
    self.assertFalse(np.allclose(out_zero_activation, 0))

class CompiledTransformerModelTest(parameterized.TestCase):

  def _get_one_hot_embed_unembed(self, vocab_size, max_seq_len):
    # Embeds tokens as one-hot into the first `vocab_size` dimensions
    token_embed = hk.Embed(
        embedding_matrix=jnp.block(
            [jnp.eye(vocab_size),
             jnp.zeros((vocab_size, max_seq_len))]))

    # Embeds positions as one-hot into the last `max_seq_len` dimensions
    position_embed = hk.Embed(
        embedding_matrix=jnp.block(
            [jnp.zeros((max_seq_len, vocab_size)),
             jnp.eye(max_seq_len)]))

    class Unembed(hk.Module):

      def __call__(self, embeddings):
        return jnp.argmax(embeddings[:, :, :vocab_size], axis=-1)

    return token_embed, position_embed, Unembed()

  def test_embedding_gives_desired_result(self):
    tokens = jnp.array([[1, 2, 3]])
    vocab_size, max_seq_len, pad_token = 5, 5, 0

    expected_embeddings = jnp.array([[[0, 1, 0, 0, 0, 1, 0, 0, 0, 0],
                                      [0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
                                      [0, 0, 0, 1, 0, 0, 0, 1, 0, 0]]])

    @hk.transform
    def embed(tokens):

      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      token_embed, position_embed, unembed = self._get_one_hot_embed_unembed(
          vocab_size, max_seq_len)
      compiled_model = model.CompiledTransformerModel(
          transformer=transformer,
          token_embed=token_embed,
          position_embed=position_embed,
          unembed=unembed,
          use_unembed_argmax=True,
          pad_token=pad_token)
      return compiled_model.embed(tokens)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params_no_mlps = self._zero_mlps(params)

    out_zero_activation = forward_zero.apply(params, next(rng), emb, mask)
    out_no_mlps = forward.apply(params_no_mlps, next(rng), emb, mask)

    self._check_layer_naming(params)
    np.testing.assert_allclose(out_zero_activation, out_no_mlps)
    self.assertFalse(np.allclose(out_zero_activation, 0))

class CompiledTransformerModelTest(parameterized.TestCase):

  def _get_one_hot_embed_unembed(self, vocab_size, max_seq_len):
    # Embeds tokens as one-hot into the first `vocab_size` dimensions
    token_embed = hk.Embed(
        embedding_matrix=jnp.block(
            [jnp.eye(vocab_size),
             jnp.zeros((vocab_size, max_seq_len))]))

    # Embeds positions as one-hot into the last `max_seq_len` dimensions
    position_embed = hk.Embed(
        embedding_matrix=jnp.block(
            [jnp.zeros((max_seq_len, vocab_size)),
             jnp.eye(max_seq_len)]))

    class Unembed(hk.Module):

      def __call__(self, embeddings):
        return jnp.argmax(embeddings[:, :, :vocab_size], axis=-1)

    return token_embed, position_embed, Unembed()

  def test_embedding_gives_desired_result(self):
    tokens = jnp.array([[1, 2, 3]])
    vocab_size, max_seq_len, pad_token = 5, 5, 0

    expected_embeddings = jnp.array([[[0, 1, 0, 0, 0, 1, 0, 0, 0, 0],
                                      [0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
                                      [0, 0, 0, 1, 0, 0, 0, 1, 0, 0]]])

    @hk.transform
    def embed(tokens):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      token_embed, position_embed, unembed = self._get_one_hot_embed_unembed(
          vocab_size, max_seq_len)
      compiled_model = model.CompiledTransformerModel(
          transformer=transformer,
          token_embed=token_embed,
          position_embed=position_embed,
          unembed=unembed,
          use_unembed_argmax=True,
          pad_token=pad_token)
      return compiled_model.embed(tokens)

    rng = hk.PRNGSequence(1)
    params = embed.init(next(rng), tokens)
    embeddings = embed.apply(params, next(rng), tokens)

    np.testing.assert_allclose(embeddings, expected_embeddings)

  def test_embedding_then_unembedding_gives_same_tokens(self):
    tokens = jnp.array([[1, 2, 3], [4, 5, 6], [3, 2, 4]])
    vocab_size, max_seq_len, pad_token = 10, 5, 0

    @hk.transform
    def embed_unembed(tokens):

      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      token_embed, position_embed, unembed = self._get_one_hot_embed_unembed(
          vocab_size, max_seq_len)
      compiled_model = model.CompiledTransformerModel(
          transformer=transformer,
          token_embed=token_embed,
          position_embed=position_embed,
          unembed=unembed,
          use_unembed_argmax=True,
          pad_token=pad_token)
      embeddings = compiled_model.embed(tokens)
      unembeddings = compiled_model.unembed(embeddings)
      return embeddings, unembeddings


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Instrumented attention layer (forked from the Haiku library implementation).
"""

from typing import Optional
import warnings

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

@chex.dataclass
class AttentionOutput:
  out: jax.Array  # [..., T', D']
  logits: jax.Array  # [..., H, T', T]

class MultiHeadAttention(hk.Module):
  """Multi-headed attention (MHA) module.

  This module is intended for attending over sequences of vectors.

  Rough sketch:
  - Compute keys (K), queries (Q), and values (V) as projections of inputs.
  - Attention weights are computed as W = softmax(QK^T / sqrt(key_size)).
  - Output is another projection of WV^T.

  For more detail, see the original Transformer paper:
    "Attention is all you need" https://arxiv.org/abs/1706.03762.

  Glossary of shapes:
  - T: Sequence length.
  - D: Vector (embedding) size.
  - H: Number of attention heads.
  """

  def __init__(
      self,
      num_heads: int,
      key_size: int,
      # TODO(b/240019186): Remove `w_init_scale`.
      w_init_scale: Optional[float] = None,
      *,
      w_init: Optional[hk.initializers.Initializer] = None,
      value_size: Optional[int] = None,
      model_size: Optional[int] = None,
      name: Optional[str] = None,
  ):
    """Initialises the module.

    Args:
      num_heads: Number of independent attention heads (H).
      key_size: The size of keys (K) and queries used for attention.
      w_init_scale: DEPRECATED. Please use w_init instead.
      w_init: Initialiser for weights in the linear map.
      value_size: Optional size of the value projection (V). If None, defaults
        to the key size (K).
      model_size: Optional size of the output embedding (D'). If None, defaults
        to the key size multiplied by the number of heads (K * H).
      name: Optional name for this module.
    """

    super().__init__(name=name)
    self.num_heads = num_heads
    self.key_size = key_size
    self.value_size = value_size or key_size
    self.model_size = model_size or key_size * num_heads

    # Backwards-compatibility for w_init_scale.
    if w_init_scale is not None:
      warnings.warn(
          "w_init_scale is deprecated; please pass an explicit weight "
          "initialiser instead.", DeprecationWarning)
    if w_init and w_init_scale:
      raise ValueError("Please provide only `w_init`, not `w_init_scale`.")
    if w_init is None and w_init_scale is None:
      raise ValueError("Please provide a weight initializer: `w_init`.")
    if w_init is None:
      w_init = hk.initializers.VarianceScaling(w_init_scale)
    self.w_init = w_init


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Instrumented attention layer (forked from the Haiku library implementation).
"""

from typing import Optional
import warnings

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

@chex.dataclass
class AttentionOutput:
  out: jax.Array  # [..., T', D']
  logits: jax.Array  # [..., H, T', T]

class MultiHeadAttention(hk.Module):
  """Multi-headed attention (MHA) module.

  This module is intended for attending over sequences of vectors.

  Rough sketch:
  - Compute keys (K), queries (Q), and values (V) as projections of inputs.
  - Attention weights are computed as W = softmax(QK^T / sqrt(key_size)).
  - Output is another projection of WV^T.

  For more detail, see the original Transformer paper:
    "Attention is all you need" https://arxiv.org/abs/1706.03762.

  Glossary of shapes:
  - T: Sequence length.
  - D: Vector (embedding) size.
  - H: Number of attention heads.
  """

  def __init__(
      self,
      num_heads: int,
      key_size: int,
      # TODO(b/240019186): Remove `w_init_scale`.
      w_init_scale: Optional[float] = None,
      *,
      w_init: Optional[hk.initializers.Initializer] = None,
      value_size: Optional[int] = None,
      model_size: Optional[int] = None,
      name: Optional[str] = None,
  ):
    """Initialises the module.

    Args:
      num_heads: Number of independent attention heads (H).
      key_size: The size of keys (K) and queries used for attention.
      w_init_scale: DEPRECATED. Please use w_init instead.
      w_init: Initialiser for weights in the linear map.
      value_size: Optional size of the value projection (V). If None, defaults
        to the key size (K).
      model_size: Optional size of the output embedding (D'). If None, defaults
        to the key size multiplied by the number of heads (K * H).
      name: Optional name for this module.
    """
    super().__init__(name=name)
    self.num_heads = num_heads
    self.key_size = key_size
    self.value_size = value_size or key_size
    self.model_size = model_size or key_size * num_heads

    # Backwards-compatibility for w_init_scale.
    if w_init_scale is not None:
      warnings.warn(
          "w_init_scale is deprecated; please pass an explicit weight "
          "initialiser instead.", DeprecationWarning)
    if w_init and w_init_scale:
      raise ValueError("Please provide only `w_init`, not `w_init_scale`.")
    if w_init is None and w_init_scale is None:
      raise ValueError("Please provide a weight initializer: `w_init`.")
    if w_init is None:
      w_init = hk.initializers.VarianceScaling(w_init_scale)
    self.w_init = w_init

  def __call__(
      self,
      query: jnp.ndarray,
      key: jnp.ndarray,
      value: jnp.ndarray,
      mask: Optional[jnp.ndarray] = None,
  ) -> AttentionOutput:
    """Computes (optionally masked) MHA with queries, keys & values.

    This module broadcasts over zero or more 'batch-like' leading dimensions.

    Args:
      query: Embeddings sequence used to compute queries; shape [..., T', D_q].
      key: Embeddings sequence used to compute keys; shape [..., T, D_k].
      value: Embeddings sequence used to compute values; shape [..., T, D_v].
      mask: Optional mask applied to attention weights; shape [..., H=1, T', T].

    Returns:
      A new sequence of embeddings, consisting of a projection of the
        attention-weighted value projections; shape [..., T', D'].
    """

    # In shape hints below, we suppress the leading dims [...] for brevity.
    # Hence e.g. [A, B] should be read in every case as [..., A, B].
    *leading_dims, sequence_length, _ = query.shape
    projection = self._linear_projection

    # Compute key/query/values (overload K/Q/V to denote the respective sizes).
    query_heads = projection(query, self.key_size, "query")  # [T', H, Q=K]
    key_heads = projection(key, self.key_size, "key")  # [T, H, K]
    value_heads = projection(value, self.value_size, "value")  # [T, H, V]

    # Compute attention weights.
    attn_logits = jnp.einsum("...thd,...Thd->...htT", query_heads, key_heads)
    attn_logits = attn_logits / np.sqrt(self.key_size).astype(key.dtype)
    if mask is not None:
      if mask.ndim != attn_logits.ndim:
        raise ValueError(
            f"Mask dimensionality {mask.ndim} must match logits dimensionality "
            f"{attn_logits.ndim}.")
      attn_logits = jnp.where(mask, attn_logits, -1e30)
    attn_weights = jax.nn.softmax(attn_logits)  # [H, T', T]

    # Weight the values by the attention and flatten the head vectors.
    attn = jnp.einsum("...htT,...Thd->...thd", attn_weights, value_heads)
    attn = jnp.reshape(attn, (*leading_dims, sequence_length, -1))  # [T', H*V]

    # Apply another projection to get the final embeddings.
    final_projection = hk.Linear(self.model_size, w_init=self.w_init)
    return AttentionOutput(
        out=final_projection(attn),
        logits=attn_logits,
    )

  @hk.transparent
  def _linear_projection(
      self,
      x: jnp.ndarray,
      head_size: int,
      name: Optional[str] = None,
  ) -> jnp.ndarray:

    y = hk.Linear(self.num_heads * head_size, w_init=self.w_init, name=name)(x)
    *leading_dims, _ = x.shape
    return y.reshape((*leading_dims, self.num_heads, head_size))


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 sys
import abc

import torch
import torch.distributed as dist

from betty.patch.data_loader import get_distributed_data_loader
from betty.patch.optimizer import patch_optimizer
from betty.patch.scheduler import patch_scheduler
from betty.configs import Config
from betty.hypergradient import get_grads
from betty.utils import convert_tensor, log_from_loss_dict

class Problem:
    """
    This is the base class for an optimization problem in multilevel optimization.
    Specifically, each problem is defined by the parameter (or module), the sets of the upper
    and lower constraining problems, the dataset, the loss function, the optimizer, and other
    optimization configurations (e.g. best-response Jacobian calculation algorithm, number of
    unrolling steps, etc.).
    """

    def __init__(
        self,
        name,
        config=None,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):
        # basic configurations
        self._name = name
        self._config = config if config is not None else Config()
        self.cfg = extra_config

        # device
        self.device = None

        # distributed
        self._strategy = None
        self.accelerator = None
        self._distributed = False
        self._backend = None
        self._world_size = None
        self._rank = None
        self._local_rank = None

        # computation graph depedency
        self._parents = []
        self._children = []
        self._paths = []

        # data loader
        self.train_data_loader = train_data_loader
        self.train_data_iterator = None
        self.cur_batch = None
        self.epoch_counter = None

        # module
        self.module = module

        # optimizer & lr scheduler
        self.optimizer = optimizer
        self.scheduler = scheduler

        # environment
        self.env = None

        # fp16 scaler
        self._fp16 = config.fp16
        self.scaler = None
        if self._fp16:
            self.initial_dynamic_scale = config.initial_dynamic_scale
            self.scale_factor = config.scale_factor

        # gradient accumulation
        self.gas = config.gradient_accumulation

        # gradient clipping
        self.gradient_clipping = config.gradient_clipping

        # warmup
        self.warmup_steps = config.warmup_steps

        # logger
        self.logger = None
        self.log_step = config.log_step
        self.log_local_step = config.log_local_step

        # step counter
        self._count = 0
        self._global_step = 0

        # misc
        self._leaf = False
        self._first_order = False
        self._retain_graph = config.retain_graph
        self._allow_unused = config.allow_unused
        self._unroll_steps = config.unroll_steps
        self._roll_back = False
        self._inner_loop_start = True
        self._training = True
        self.ready = None

    def initialize(self):
        """
        ``initialize`` patches/sets up module, optimizer, data loader, etc. after compiling a
        user-provided configuration (e.g., fp16 training, iterative differentiation)
        """
        # initialize update ready to False
        self.ready = [False for _ in range(len(self._children))]

        # compile parents configurations
        first_order = []
        for problem in self._parents:
            parent_config = problem.config
            first_order.append(parent_config.first_order)
        self._first_order = all(first_order)

        # set inner_loop_start to True
        self._inner_loop_start = True

        # accelerate
        if self._strategy == "accelerate":
            from accelerate import Accelerator

            self.accelerator = Accelerator()

        # set up data loader
        if self.is_implemented("configure_train_data_loader"):
            if self.train_data_loader is None:
                self.train_data_loader = self.configure_train_data_loader()
        if self.train_data_loader is not None:
            if not isinstance(self.train_data_loader, tuple):
                self.train_data_loader = (self.train_data_loader,)
        else:
            assert self.is_implemented("get_batch")

        # set up module
        if self.is_implemented("configure_module"):
            if self.module is None:
                self.module = self.configure_module()
        assert self.module is not None, "Module must be specified!"

        # set up optimizer
        if self.is_implemented("configure_optimizer"):
            if self.optimizer is None:
                self.optimizer = self.configure_optimizer()

        # set up lr scheduler
        if self.is_implemented("configure_scheduler"):
            if self.scheduler is None:
                self.scheduler = self.configure_scheduler()

        # set up fp16 training
        if self._is_default_fp16():
            assert torch.cuda.is_available()
            scaler_cls = torch.cuda.amp.GradScaler
            if self._strategy == "fsdp":
                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

                scaler_cls = ShardedGradScaler
            self.scaler = scaler_cls(
                init_scale=self.initial_dynamic_scale, growth_factor=self.scale_factor
            )

        # patch module, optimizer, data loader, and scheduler
        self.patch_everything()

        # make train_data_loader as iterator
        if self.train_data_loader is not None:
            self.train_data_iterator = []
            self.epoch_counter = []
            for train_data_loader in self.train_data_loader:
                self.train_data_iterator.append(iter(train_data_loader))
                self.epoch_counter.append(0)

        # Logging INFO
        path_str = [[node.name for node in path] for path in self._paths]
        children_str = [node.name for node in self._children]
        parents_str = [node.name for node in self._parents]
        if self.is_rank_zero():
            self.logger.info("*** Problem Information ***")
            self.logger.info(f"Name: {self._name}")
            self.logger.info(f"Uppers: {parents_str}")
            self.logger.info(f"Lowers: {children_str}")
            self.logger.info(f"Paths: {path_str}\n")

    def patch_everything(self):
        """
        We patch module, optimizer, data loader, and lr scheduler for device placement,
        distributed training, zero optimizer, fsdp, etc.
        """

        self.patch_module()
        self.patch_optimizer()
        if self.scheduler is not None:
            self.patch_scheduler()
        if self.train_data_loader is not None:
            self.train_data_loader = [
                self.patch_data_loader(data_loader)
                for data_loader in self.train_data_loader
            ]


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 sys
import abc

import torch
import torch.distributed as dist

from betty.patch.data_loader import get_distributed_data_loader
from betty.patch.optimizer import patch_optimizer
from betty.patch.scheduler import patch_scheduler
from betty.configs import Config
from betty.hypergradient import get_grads
from betty.utils import convert_tensor, log_from_loss_dict

class Problem:
    """
    This is the base class for an optimization problem in multilevel optimization.
    Specifically, each problem is defined by the parameter (or module), the sets of the upper
    and lower constraining problems, the dataset, the loss function, the optimizer, and other
    optimization configurations (e.g. best-response Jacobian calculation algorithm, number of
    unrolling steps, etc.).
    """

    def __init__(
        self,
        name,
        config=None,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):
        # basic configurations
        self._name = name
        self._config = config if config is not None else Config()
        self.cfg = extra_config

        # device
        self.device = None

        # distributed
        self._strategy = None
        self.accelerator = None
        self._distributed = False
        self._backend = None
        self._world_size = None
        self._rank = None
        self._local_rank = None

        # computation graph depedency
        self._parents = []
        self._children = []
        self._paths = []

        # data loader
        self.train_data_loader = train_data_loader
        self.train_data_iterator = None
        self.cur_batch = None
        self.epoch_counter = None

        # module
        self.module = module

        # optimizer & lr scheduler
        self.optimizer = optimizer
        self.scheduler = scheduler

        # environment
        self.env = None

        # fp16 scaler
        self._fp16 = config.fp16
        self.scaler = None
        if self._fp16:
            self.initial_dynamic_scale = config.initial_dynamic_scale
            self.scale_factor = config.scale_factor

        # gradient accumulation
        self.gas = config.gradient_accumulation

        # gradient clipping
        self.gradient_clipping = config.gradient_clipping

        # warmup
        self.warmup_steps = config.warmup_steps

        # logger
        self.logger = None
        self.log_step = config.log_step
        self.log_local_step = config.log_local_step

        # step counter
        self._count = 0
        self._global_step = 0

        # misc
        self._leaf = False
        self._first_order = False
        self._retain_graph = config.retain_graph
        self._allow_unused = config.allow_unused
        self._unroll_steps = config.unroll_steps
        self._roll_back = False
        self._inner_loop_start = True
        self._training = True
        self.ready = None

    def initialize(self):
        """
        ``initialize`` patches/sets up module, optimizer, data loader, etc. after compiling a
        user-provided configuration (e.g., fp16 training, iterative differentiation)
        """
        # initialize update ready to False
        self.ready = [False for _ in range(len(self._children))]

        # compile parents configurations
        first_order = []
        for problem in self._parents:
            parent_config = problem.config
            first_order.append(parent_config.first_order)
        self._first_order = all(first_order)

        # set inner_loop_start to True
        self._inner_loop_start = True

        # accelerate
        if self._strategy == "accelerate":
            from accelerate import Accelerator

            self.accelerator = Accelerator()

        # set up data loader
        if self.is_implemented("configure_train_data_loader"):
            if self.train_data_loader is None:
                self.train_data_loader = self.configure_train_data_loader()
        if self.train_data_loader is not None:
            if not isinstance(self.train_data_loader, tuple):
                self.train_data_loader = (self.train_data_loader,)
        else:
            assert self.is_implemented("get_batch")

        # set up module
        if self.is_implemented("configure_module"):
            if self.module is None:
                self.module = self.configure_module()
        assert self.module is not None, "Module must be specified!"

        # set up optimizer
        if self.is_implemented("configure_optimizer"):
            if self.optimizer is None:
                self.optimizer = self.configure_optimizer()

        # set up lr scheduler
        if self.is_implemented("configure_scheduler"):
            if self.scheduler is None:
                self.scheduler = self.configure_scheduler()

        # set up fp16 training
        if self._is_default_fp16():
            assert torch.cuda.is_available()
            scaler_cls = torch.cuda.amp.GradScaler
            if self._strategy == "fsdp":
                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

                scaler_cls = ShardedGradScaler
            self.scaler = scaler_cls(
                init_scale=self.initial_dynamic_scale, growth_factor=self.scale_factor
            )

        # patch module, optimizer, data loader, and scheduler
        self.patch_everything()

        # make train_data_loader as iterator
        if self.train_data_loader is not None:
            self.train_data_iterator = []
            self.epoch_counter = []
            for train_data_loader in self.train_data_loader:
                self.train_data_iterator.append(iter(train_data_loader))
                self.epoch_counter.append(0)

        # Logging INFO
        path_str = [[node.name for node in path] for path in self._paths]
        children_str = [node.name for node in self._children]
        parents_str = [node.name for node in self._parents]
        if self.is_rank_zero():
            self.logger.info("*** Problem Information ***")
            self.logger.info(f"Name: {self._name}")
            self.logger.info(f"Uppers: {parents_str}")
            self.logger.info(f"Lowers: {children_str}")
            self.logger.info(f"Paths: {path_str}\n")

    def patch_everything(self):
        """
        We patch module, optimizer, data loader, and lr scheduler for device placement,
        distributed training, zero optimizer, fsdp, etc.
        """
        self.patch_module()
        self.patch_optimizer()
        if self.scheduler is not None:
            self.patch_scheduler()
        if self.train_data_loader is not None:
            self.train_data_loader = [
                self.patch_data_loader(data_loader)
                for data_loader in self.train_data_loader
            ]

    def patch_module(self):
        """
        Patch module given the systems configuration (e.g., DDP, FSDP)
        """
        self.module.to(self.device)
        if self._strategy in ["distributed", "zero"]:
            self.synchronize_params(self.parameters())
            self.module = torch.nn.parallel.DistributedDataParallel(
                module=self.module,
                gradient_as_bucket_view=True,
            )
        elif self._strategy == "fsdp":
            if self.is_rank_zero():
                self.logger.warning("FSDP requires PyTorch version >= 1.12")
            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

            self.synchronize_params(self.parameters())
            self.module = FSDP(self.module, device_id=self.device)
        elif self._strategy == "accelerate":
            self.module = self.accelerator.prepare(self.module)

    def patch_optimizer(self):
        """
        Patch optimizer given the systems configuration (e.g., DDP, FSDP)
        """

        params = self.trainable_parameters()
        if self.is_implemented("param_groups") and self._strategy != "fsdp":
            params = self.param_groups()
        is_zero = True if self._strategy == "zero" else False
        if self._strategy == "accelerate":
            self.optimizer = self.accelerator.prepare(self.optimizer)
        else:
            self.optimizer = patch_optimizer(self.optimizer, params, is_zero)


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 sys
import abc

import torch
import torch.distributed as dist

from betty.patch.data_loader import get_distributed_data_loader
from betty.patch.optimizer import patch_optimizer
from betty.patch.scheduler import patch_scheduler
from betty.configs import Config
from betty.hypergradient import get_grads
from betty.utils import convert_tensor, log_from_loss_dict

class Problem:
    """
    This is the base class for an optimization problem in multilevel optimization.
    Specifically, each problem is defined by the parameter (or module), the sets of the upper
    and lower constraining problems, the dataset, the loss function, the optimizer, and other
    optimization configurations (e.g. best-response Jacobian calculation algorithm, number of
    unrolling steps, etc.).
    """

    def __init__(
        self,
        name,
        config=None,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):
        # basic configurations
        self._name = name
        self._config = config if config is not None else Config()
        self.cfg = extra_config

        # device
        self.device = None

        # distributed
        self._strategy = None
        self.accelerator = None
        self._distributed = False
        self._backend = None
        self._world_size = None
        self._rank = None
        self._local_rank = None

        # computation graph depedency
        self._parents = []
        self._children = []
        self._paths = []

        # data loader
        self.train_data_loader = train_data_loader
        self.train_data_iterator = None
        self.cur_batch = None
        self.epoch_counter = None

        # module
        self.module = module

        # optimizer & lr scheduler
        self.optimizer = optimizer
        self.scheduler = scheduler

        # environment
        self.env = None

        # fp16 scaler
        self._fp16 = config.fp16
        self.scaler = None
        if self._fp16:
            self.initial_dynamic_scale = config.initial_dynamic_scale
            self.scale_factor = config.scale_factor

        # gradient accumulation
        self.gas = config.gradient_accumulation

        # gradient clipping
        self.gradient_clipping = config.gradient_clipping

        # warmup
        self.warmup_steps = config.warmup_steps

        # logger
        self.logger = None
        self.log_step = config.log_step
        self.log_local_step = config.log_local_step

        # step counter
        self._count = 0
        self._global_step = 0

        # misc
        self._leaf = False
        self._first_order = False
        self._retain_graph = config.retain_graph
        self._allow_unused = config.allow_unused
        self._unroll_steps = config.unroll_steps
        self._roll_back = False
        self._inner_loop_start = True
        self._training = True
        self.ready = None

    def initialize(self):
        """
        ``initialize`` patches/sets up module, optimizer, data loader, etc. after compiling a
        user-provided configuration (e.g., fp16 training, iterative differentiation)
        """
        # initialize update ready to False
        self.ready = [False for _ in range(len(self._children))]

        # compile parents configurations
        first_order = []
        for problem in self._parents:
            parent_config = problem.config
            first_order.append(parent_config.first_order)
        self._first_order = all(first_order)

        # set inner_loop_start to True
        self._inner_loop_start = True

        # accelerate
        if self._strategy == "accelerate":
            from accelerate import Accelerator

            self.accelerator = Accelerator()

        # set up data loader
        if self.is_implemented("configure_train_data_loader"):
            if self.train_data_loader is None:
                self.train_data_loader = self.configure_train_data_loader()
        if self.train_data_loader is not None:
            if not isinstance(self.train_data_loader, tuple):
                self.train_data_loader = (self.train_data_loader,)
        else:
            assert self.is_implemented("get_batch")

        # set up module
        if self.is_implemented("configure_module"):
            if self.module is None:
                self.module = self.configure_module()
        assert self.module is not None, "Module must be specified!"

        # set up optimizer
        if self.is_implemented("configure_optimizer"):
            if self.optimizer is None:
                self.optimizer = self.configure_optimizer()

        # set up lr scheduler
        if self.is_implemented("configure_scheduler"):
            if self.scheduler is None:
                self.scheduler = self.configure_scheduler()

        # set up fp16 training
        if self._is_default_fp16():
            assert torch.cuda.is_available()
            scaler_cls = torch.cuda.amp.GradScaler
            if self._strategy == "fsdp":
                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

                scaler_cls = ShardedGradScaler
            self.scaler = scaler_cls(
                init_scale=self.initial_dynamic_scale, growth_factor=self.scale_factor
            )

        # patch module, optimizer, data loader, and scheduler
        self.patch_everything()

        # make train_data_loader as iterator
        if self.train_data_loader is not None:
            self.train_data_iterator = []
            self.epoch_counter = []
            for train_data_loader in self.train_data_loader:
                self.train_data_iterator.append(iter(train_data_loader))
                self.epoch_counter.append(0)

        # Logging INFO
        path_str = [[node.name for node in path] for path in self._paths]
        children_str = [node.name for node in self._children]
        parents_str = [node.name for node in self._parents]
        if self.is_rank_zero():
            self.logger.info("*** Problem Information ***")
            self.logger.info(f"Name: {self._name}")
            self.logger.info(f"Uppers: {parents_str}")
            self.logger.info(f"Lowers: {children_str}")
            self.logger.info(f"Paths: {path_str}\n")

    def patch_everything(self):
        """
        We patch module, optimizer, data loader, and lr scheduler for device placement,
        distributed training, zero optimizer, fsdp, etc.
        """
        self.patch_module()
        self.patch_optimizer()
        if self.scheduler is not None:
            self.patch_scheduler()
        if self.train_data_loader is not None:
            self.train_data_loader = [
                self.patch_data_loader(data_loader)
                for data_loader in self.train_data_loader
            ]

    def patch_module(self):
        """
        Patch module given the systems configuration (e.g., DDP, FSDP)
        """
        self.module.to(self.device)
        if self._strategy in ["distributed", "zero"]:
            self.synchronize_params(self.parameters())
            self.module = torch.nn.parallel.DistributedDataParallel(
                module=self.module,
                gradient_as_bucket_view=True,
            )
        elif self._strategy == "fsdp":
            if self.is_rank_zero():
                self.logger.warning("FSDP requires PyTorch version >= 1.12")
            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

            self.synchronize_params(self.parameters())
            self.module = FSDP(self.module, device_id=self.device)
        elif self._strategy == "accelerate":
            self.module = self.accelerator.prepare(self.module)

    def patch_optimizer(self):
        """
        Patch optimizer given the systems configuration (e.g., DDP, FSDP)
        """
        params = self.trainable_parameters()
        if self.is_implemented("param_groups") and self._strategy != "fsdp":
            params = self.param_groups()
        is_zero = True if self._strategy == "zero" else False
        if self._strategy == "accelerate":
            self.optimizer = self.accelerator.prepare(self.optimizer)
        else:
            self.optimizer = patch_optimizer(self.optimizer, params, is_zero)

    def patch_scheduler(self):
        """
        Patch scheduler given the systems configuration (e.g., DDP, FSDP)
        """
        self.scheduler = patch_scheduler(self.scheduler, self.optimizer)
        if self._strategy == "accelerate":
            self.scheduler = self.accelerator.prepare(self.scheduler)

    def patch_data_loader(self, loader):
        """
        Patch data loader given the systems configuration (e.g., DDP, FSDP)
        """

        if self._strategy in ["distributed", "zero", "fsdp"]:
            patched_loader = get_distributed_data_loader(
                loader, world_size=self._world_size, rank=self._rank
            )
        elif self._strategy == "accelerate":
            patched_loader = self.accelerator.prepare(loader)
        else:
            patched_loader = loader

        return patched_loader


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 sys
import abc

import torch
import torch.distributed as dist

from betty.patch.data_loader import get_distributed_data_loader
from betty.patch.optimizer import patch_optimizer
from betty.patch.scheduler import patch_scheduler
from betty.configs import Config
from betty.hypergradient import get_grads
from betty.utils import convert_tensor, log_from_loss_dict

class Problem:
    """
    This is the base class for an optimization problem in multilevel optimization.
    Specifically, each problem is defined by the parameter (or module), the sets of the upper
    and lower constraining problems, the dataset, the loss function, the optimizer, and other
    optimization configurations (e.g. best-response Jacobian calculation algorithm, number of
    unrolling steps, etc.).
    """

    def __init__(
        self,
        name,
        config=None,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):
        # basic configurations
        self._name = name
        self._config = config if config is not None else Config()
        self.cfg = extra_config

        # device
        self.device = None

        # distributed
        self._strategy = None
        self.accelerator = None
        self._distributed = False
        self._backend = None
        self._world_size = None
        self._rank = None
        self._local_rank = None

        # computation graph depedency
        self._parents = []
        self._children = []
        self._paths = []

        # data loader
        self.train_data_loader = train_data_loader
        self.train_data_iterator = None
        self.cur_batch = None
        self.epoch_counter = None

        # module
        self.module = module

        # optimizer & lr scheduler
        self.optimizer = optimizer
        self.scheduler = scheduler

        # environment
        self.env = None

        # fp16 scaler
        self._fp16 = config.fp16
        self.scaler = None
        if self._fp16:
            self.initial_dynamic_scale = config.initial_dynamic_scale
            self.scale_factor = config.scale_factor

        # gradient accumulation
        self.gas = config.gradient_accumulation

        # gradient clipping
        self.gradient_clipping = config.gradient_clipping

        # warmup
        self.warmup_steps = config.warmup_steps

        # logger
        self.logger = None
        self.log_step = config.log_step
        self.log_local_step = config.log_local_step

        # step counter
        self._count = 0
        self._global_step = 0

        # misc
        self._leaf = False
        self._first_order = False
        self._retain_graph = config.retain_graph
        self._allow_unused = config.allow_unused
        self._unroll_steps = config.unroll_steps
        self._roll_back = False
        self._inner_loop_start = True
        self._training = True
        self.ready = None

    def initialize(self):
        """
        ``initialize`` patches/sets up module, optimizer, data loader, etc. after compiling a
        user-provided configuration (e.g., fp16 training, iterative differentiation)
        """
        # initialize update ready to False
        self.ready = [False for _ in range(len(self._children))]

        # compile parents configurations
        first_order = []
        for problem in self._parents:
            parent_config = problem.config
            first_order.append(parent_config.first_order)
        self._first_order = all(first_order)

        # set inner_loop_start to True
        self._inner_loop_start = True

        # accelerate
        if self._strategy == "accelerate":
            from accelerate import Accelerator

            self.accelerator = Accelerator()

        # set up data loader
        if self.is_implemented("configure_train_data_loader"):
            if self.train_data_loader is None:
                self.train_data_loader = self.configure_train_data_loader()
        if self.train_data_loader is not None:
            if not isinstance(self.train_data_loader, tuple):
                self.train_data_loader = (self.train_data_loader,)
        else:
            assert self.is_implemented("get_batch")

        # set up module
        if self.is_implemented("configure_module"):
            if self.module is None:
                self.module = self.configure_module()
        assert self.module is not None, "Module must be specified!"

        # set up optimizer
        if self.is_implemented("configure_optimizer"):
            if self.optimizer is None:
                self.optimizer = self.configure_optimizer()

        # set up lr scheduler
        if self.is_implemented("configure_scheduler"):
            if self.scheduler is None:
                self.scheduler = self.configure_scheduler()

        # set up fp16 training
        if self._is_default_fp16():
            assert torch.cuda.is_available()
            scaler_cls = torch.cuda.amp.GradScaler
            if self._strategy == "fsdp":
                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

                scaler_cls = ShardedGradScaler
            self.scaler = scaler_cls(
                init_scale=self.initial_dynamic_scale, growth_factor=self.scale_factor
            )

        # patch module, optimizer, data loader, and scheduler
        self.patch_everything()

        # make train_data_loader as iterator
        if self.train_data_loader is not None:
            self.train_data_iterator = []
            self.epoch_counter = []
            for train_data_loader in self.train_data_loader:
                self.train_data_iterator.append(iter(train_data_loader))
                self.epoch_counter.append(0)

        # Logging INFO
        path_str = [[node.name for node in path] for path in self._paths]
        children_str = [node.name for node in self._children]
        parents_str = [node.name for node in self._parents]
        if self.is_rank_zero():
            self.logger.info("*** Problem Information ***")
            self.logger.info(f"Name: {self._name}")
            self.logger.info(f"Uppers: {parents_str}")
            self.logger.info(f"Lowers: {children_str}")
            self.logger.info(f"Paths: {path_str}\n")

    def patch_everything(self):
        """
        We patch module, optimizer, data loader, and lr scheduler for device placement,
        distributed training, zero optimizer, fsdp, etc.
        """
        self.patch_module()
        self.patch_optimizer()
        if self.scheduler is not None:
            self.patch_scheduler()
        if self.train_data_loader is not None:
            self.train_data_loader = [
                self.patch_data_loader(data_loader)
                for data_loader in self.train_data_loader
            ]

    def patch_module(self):
        """
        Patch module given the systems configuration (e.g., DDP, FSDP)
        """
        self.module.to(self.device)
        if self._strategy in ["distributed", "zero"]:
            self.synchronize_params(self.parameters())
            self.module = torch.nn.parallel.DistributedDataParallel(
                module=self.module,
                gradient_as_bucket_view=True,
            )
        elif self._strategy == "fsdp":
            if self.is_rank_zero():
                self.logger.warning("FSDP requires PyTorch version >= 1.12")
            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

            self.synchronize_params(self.parameters())
            self.module = FSDP(self.module, device_id=self.device)
        elif self._strategy == "accelerate":
            self.module = self.accelerator.prepare(self.module)

    def patch_optimizer(self):
        """
        Patch optimizer given the systems configuration (e.g., DDP, FSDP)
        """
        params = self.trainable_parameters()
        if self.is_implemented("param_groups") and self._strategy != "fsdp":
            params = self.param_groups()
        is_zero = True if self._strategy == "zero" else False
        if self._strategy == "accelerate":
            self.optimizer = self.accelerator.prepare(self.optimizer)
        else:
            self.optimizer = patch_optimizer(self.optimizer, params, is_zero)

    def patch_scheduler(self):
        """
        Patch scheduler given the systems configuration (e.g., DDP, FSDP)
        """
        self.scheduler = patch_scheduler(self.scheduler, self.optimizer)
        if self._strategy == "accelerate":
            self.scheduler = self.accelerator.prepare(self.scheduler)

    def patch_data_loader(self, loader):
        """
        Patch data loader given the systems configuration (e.g., DDP, FSDP)
        """
        if self._strategy in ["distributed", "zero", "fsdp"]:
            patched_loader = get_distributed_data_loader(
                loader, world_size=self._world_size, rank=self._rank
            )
        elif self._strategy == "accelerate":
            patched_loader = self.accelerator.prepare(loader)
        else:
            patched_loader = loader

        return patched_loader

    def set_module(self, module):
        """
        Set new module for the current Problem class.
        """
        self.module = module
        self.patch_module()

    def set_optimizer(self, optimizer):
        """
        Set new optimizer for the current Problem class.
        """
        self.optimizer = optimizer
        self.patch_optimizer()

    def set_scheduler(self, scheduler):
        """
        Set new scheduler for the current Problem class.
        """
        self.scheduler = scheduler
        self.patch_scheduler()

    def set_train_data_loader(self, loader, idx=0):
        """
        Set new data loader for the current Problem class.
        """
        self.train_data_loader[idx] = self.patch_data_loader(loader)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def forward(self, *args, **kwargs):
        """
        Users define how forward (or call) function is defined for the problem here.
        """
        return self.module(*args, **kwargs)

    @abc.abstractmethod
    def training_step(self, batch):
        """
        Users define the loss function of the problem here.
        """
        raise NotImplementedError

    def training_step_exec(self, batch):

        if self._is_default_fp16():
            with torch.cuda.amp.autocast():
                return self.training_step(batch)
        else:
            return self.training_step(batch)


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 sys
import abc

import torch
import torch.distributed as dist

from betty.patch.data_loader import get_distributed_data_loader
from betty.patch.optimizer import patch_optimizer
from betty.patch.scheduler import patch_scheduler
from betty.configs import Config
from betty.hypergradient import get_grads
from betty.utils import convert_tensor, log_from_loss_dict

class Problem:
    """
    This is the base class for an optimization problem in multilevel optimization.
    Specifically, each problem is defined by the parameter (or module), the sets of the upper
    and lower constraining problems, the dataset, the loss function, the optimizer, and other
    optimization configurations (e.g. best-response Jacobian calculation algorithm, number of
    unrolling steps, etc.).
    """

    def __init__(
        self,
        name,
        config=None,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):
        # basic configurations
        self._name = name
        self._config = config if config is not None else Config()
        self.cfg = extra_config

        # device
        self.device = None

        # distributed
        self._strategy = None
        self.accelerator = None
        self._distributed = False
        self._backend = None
        self._world_size = None
        self._rank = None
        self._local_rank = None

        # computation graph depedency
        self._parents = []
        self._children = []
        self._paths = []

        # data loader
        self.train_data_loader = train_data_loader
        self.train_data_iterator = None
        self.cur_batch = None
        self.epoch_counter = None

        # module
        self.module = module

        # optimizer & lr scheduler
        self.optimizer = optimizer
        self.scheduler = scheduler

        # environment
        self.env = None

        # fp16 scaler
        self._fp16 = config.fp16
        self.scaler = None
        if self._fp16:
            self.initial_dynamic_scale = config.initial_dynamic_scale
            self.scale_factor = config.scale_factor

        # gradient accumulation
        self.gas = config.gradient_accumulation

        # gradient clipping
        self.gradient_clipping = config.gradient_clipping

        # warmup
        self.warmup_steps = config.warmup_steps

        # logger
        self.logger = None
        self.log_step = config.log_step
        self.log_local_step = config.log_local_step

        # step counter
        self._count = 0
        self._global_step = 0

        # misc
        self._leaf = False
        self._first_order = False
        self._retain_graph = config.retain_graph
        self._allow_unused = config.allow_unused
        self._unroll_steps = config.unroll_steps
        self._roll_back = False
        self._inner_loop_start = True
        self._training = True
        self.ready = None

    def initialize(self):
        """
        ``initialize`` patches/sets up module, optimizer, data loader, etc. after compiling a
        user-provided configuration (e.g., fp16 training, iterative differentiation)
        """
        # initialize update ready to False
        self.ready = [False for _ in range(len(self._children))]

        # compile parents configurations
        first_order = []
        for problem in self._parents:
            parent_config = problem.config
            first_order.append(parent_config.first_order)
        self._first_order = all(first_order)

        # set inner_loop_start to True
        self._inner_loop_start = True

        # accelerate
        if self._strategy == "accelerate":
            from accelerate import Accelerator

            self.accelerator = Accelerator()

        # set up data loader
        if self.is_implemented("configure_train_data_loader"):
            if self.train_data_loader is None:
                self.train_data_loader = self.configure_train_data_loader()
        if self.train_data_loader is not None:
            if not isinstance(self.train_data_loader, tuple):
                self.train_data_loader = (self.train_data_loader,)
        else:
            assert self.is_implemented("get_batch")

        # set up module
        if self.is_implemented("configure_module"):
            if self.module is None:
                self.module = self.configure_module()
        assert self.module is not None, "Module must be specified!"

        # set up optimizer
        if self.is_implemented("configure_optimizer"):
            if self.optimizer is None:
                self.optimizer = self.configure_optimizer()

        # set up lr scheduler
        if self.is_implemented("configure_scheduler"):
            if self.scheduler is None:
                self.scheduler = self.configure_scheduler()

        # set up fp16 training
        if self._is_default_fp16():
            assert torch.cuda.is_available()
            scaler_cls = torch.cuda.amp.GradScaler
            if self._strategy == "fsdp":
                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

                scaler_cls = ShardedGradScaler
            self.scaler = scaler_cls(
                init_scale=self.initial_dynamic_scale, growth_factor=self.scale_factor
            )

        # patch module, optimizer, data loader, and scheduler
        self.patch_everything()

        # make train_data_loader as iterator
        if self.train_data_loader is not None:
            self.train_data_iterator = []
            self.epoch_counter = []
            for train_data_loader in self.train_data_loader:
                self.train_data_iterator.append(iter(train_data_loader))
                self.epoch_counter.append(0)

        # Logging INFO
        path_str = [[node.name for node in path] for path in self._paths]
        children_str = [node.name for node in self._children]
        parents_str = [node.name for node in self._parents]
        if self.is_rank_zero():
            self.logger.info("*** Problem Information ***")
            self.logger.info(f"Name: {self._name}")
            self.logger.info(f"Uppers: {parents_str}")
            self.logger.info(f"Lowers: {children_str}")
            self.logger.info(f"Paths: {path_str}\n")

    def patch_everything(self):
        """
        We patch module, optimizer, data loader, and lr scheduler for device placement,
        distributed training, zero optimizer, fsdp, etc.
        """
        self.patch_module()
        self.patch_optimizer()
        if self.scheduler is not None:
            self.patch_scheduler()
        if self.train_data_loader is not None:
            self.train_data_loader = [
                self.patch_data_loader(data_loader)
                for data_loader in self.train_data_loader
            ]

    def patch_module(self):
        """
        Patch module given the systems configuration (e.g., DDP, FSDP)
        """
        self.module.to(self.device)
        if self._strategy in ["distributed", "zero"]:
            self.synchronize_params(self.parameters())
            self.module = torch.nn.parallel.DistributedDataParallel(
                module=self.module,
                gradient_as_bucket_view=True,
            )
        elif self._strategy == "fsdp":
            if self.is_rank_zero():
                self.logger.warning("FSDP requires PyTorch version >= 1.12")
            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

            self.synchronize_params(self.parameters())
            self.module = FSDP(self.module, device_id=self.device)
        elif self._strategy == "accelerate":
            self.module = self.accelerator.prepare(self.module)

    def patch_optimizer(self):
        """
        Patch optimizer given the systems configuration (e.g., DDP, FSDP)
        """
        params = self.trainable_parameters()
        if self.is_implemented("param_groups") and self._strategy != "fsdp":
            params = self.param_groups()
        is_zero = True if self._strategy == "zero" else False
        if self._strategy == "accelerate":
            self.optimizer = self.accelerator.prepare(self.optimizer)
        else:
            self.optimizer = patch_optimizer(self.optimizer, params, is_zero)

    def patch_scheduler(self):
        """
        Patch scheduler given the systems configuration (e.g., DDP, FSDP)
        """
        self.scheduler = patch_scheduler(self.scheduler, self.optimizer)
        if self._strategy == "accelerate":
            self.scheduler = self.accelerator.prepare(self.scheduler)

    def patch_data_loader(self, loader):
        """
        Patch data loader given the systems configuration (e.g., DDP, FSDP)
        """
        if self._strategy in ["distributed", "zero", "fsdp"]:
            patched_loader = get_distributed_data_loader(
                loader, world_size=self._world_size, rank=self._rank
            )
        elif self._strategy == "accelerate":
            patched_loader = self.accelerator.prepare(loader)
        else:
            patched_loader = loader

        return patched_loader

    def set_module(self, module):
        """
        Set new module for the current Problem class.
        """
        self.module = module
        self.patch_module()

    def set_optimizer(self, optimizer):
        """
        Set new optimizer for the current Problem class.
        """
        self.optimizer = optimizer
        self.patch_optimizer()

    def set_scheduler(self, scheduler):
        """
        Set new scheduler for the current Problem class.
        """
        self.scheduler = scheduler
        self.patch_scheduler()

    def set_train_data_loader(self, loader, idx=0):
        """
        Set new data loader for the current Problem class.
        """
        self.train_data_loader[idx] = self.patch_data_loader(loader)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def forward(self, *args, **kwargs):
        """
        Users define how forward (or call) function is defined for the problem here.
        """
        return self.module(*args, **kwargs)

    @abc.abstractmethod
    def training_step(self, batch):
        """
        Users define the loss function of the problem here.
        """
        raise NotImplementedError

    def training_step_exec(self, batch):
        if self._is_default_fp16():
            with torch.cuda.amp.autocast():
                return self.training_step(batch)
        else:
            return self.training_step(batch)

    def one_step_descent(self, batch=None):
        # load data
        if batch is None:
            self.cur_batch = self.get_batch()
            batch = self.cur_batch

        # calculate loss
        loss, loss_dict = self.get_loss(batch)

        # calculate gradient (a.k.a backward)
        self.backward(
            loss=loss,
            params=self.trainable_parameters(),
            paths=self._paths,
            create_graph=not self._first_order,
            retain_graph=self._retain_graph,
            allow_unused=self._allow_unused,
        )
        if self.is_implemented("grad_callback"):
            self.grad_callback()

        # calculate parameter update
        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()


Prompt:
.scheduler import patch_scheduler
from betty.configs import Config
from betty.hypergradient import get_grads
from betty.utils import convert_tensor, log_from_loss_dict

class Problem:
    """
    This is the base class for an optimization problem in multilevel optimization.
    Specifically, each problem is defined by the parameter (or module), the sets of the upper
    and lower constraining problems, the dataset, the loss function, the optimizer, and other
    optimization configurations (e.g. best-response Jacobian calculation algorithm, number of
    unrolling steps, etc.).
    """

    def __init__(
        self,
        name,
        config=None,
        module=None,
        optimizer=None,
        scheduler=None,
        train_data_loader=None,
        extra_config=None,
    ):
        # basic configurations
        self._name = name
        self._config = config if config is not None else Config()
        self.cfg = extra_config

        # device
        self.device = None

        # distributed
        self._strategy = None
        self.accelerator = None
        self._distributed = False
        self._backend = None
        self._world_size = None
        self._rank = None
        self._local_rank = None

        # computation graph depedency
        self._parents = []
        self._children = []
        self._paths = []

        # data loader
        self.train_data_loader = train_data_loader
        self.train_data_iterator = None
        self.cur_batch = None
        self.epoch_counter = None

        # module
        self.module = module

        # optimizer & lr scheduler
        self.optimizer = optimizer
        self.scheduler = scheduler

        # environment
        self.env = None

        # fp16 scaler
        self._fp16 = config.fp16
        self.scaler = None
        if self._fp16:
            self.initial_dynamic_scale = config.initial_dynamic_scale
            self.scale_factor = config.scale_factor

        # gradient accumulation
        self.gas = config.gradient_accumulation

        # gradient clipping
        self.gradient_clipping = config.gradient_clipping

        # warmup
        self.warmup_steps = config.warmup_steps

        # logger
        self.logger = None
        self.log_step = config.log_step
        self.log_local_step = config.log_local_step

        # step counter
        self._count = 0
        self._global_step = 0

        # misc
        self._leaf = False
        self._first_order = False
        self._retain_graph = config.retain_graph
        self._allow_unused = config.allow_unused
        self._unroll_steps = config.unroll_steps
        self._roll_back = False
        self._inner_loop_start = True
        self._training = True
        self.ready = None

    def initialize(self):
        """
        ``initialize`` patches/sets up module, optimizer, data loader, etc. after compiling a
        user-provided configuration (e.g., fp16 training, iterative differentiation)
        """
        # initialize update ready to False
        self.ready = [False for _ in range(len(self._children))]

        # compile parents configurations
        first_order = []
        for problem in self._parents:
            parent_config = problem.config
            first_order.append(parent_config.first_order)
        self._first_order = all(first_order)

        # set inner_loop_start to True
        self._inner_loop_start = True

        # accelerate
        if self._strategy == "accelerate":
            from accelerate import Accelerator

            self.accelerator = Accelerator()

        # set up data loader
        if self.is_implemented("configure_train_data_loader"):
            if self.train_data_loader is None:
                self.train_data_loader = self.configure_train_data_loader()
        if self.train_data_loader is not None:
            if not isinstance(self.train_data_loader, tuple):
                self.train_data_loader = (self.train_data_loader,)
        else:
            assert self.is_implemented("get_batch")

        # set up module
        if self.is_implemented("configure_module"):
            if self.module is None:
                self.module = self.configure_module()
        assert self.module is not None, "Module must be specified!"

        # set up optimizer
        if self.is_implemented("configure_optimizer"):
            if self.optimizer is None:
                self.optimizer = self.configure_optimizer()

        # set up lr scheduler
        if self.is_implemented("configure_scheduler"):
            if self.scheduler is None:
                self.scheduler = self.configure_scheduler()

        # set up fp16 training
        if self._is_default_fp16():
            assert torch.cuda.is_available()
            scaler_cls = torch.cuda.amp.GradScaler
            if self._strategy == "fsdp":
                from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

                scaler_cls = ShardedGradScaler
            self.scaler = scaler_cls(
                init_scale=self.initial_dynamic_scale, growth_factor=self.scale_factor
            )

        # patch module, optimizer, data loader, and scheduler
        self.patch_everything()

        # make train_data_loader as iterator
        if self.train_data_loader is not None:
            self.train_data_iterator = []
            self.epoch_counter = []
            for train_data_loader in self.train_data_loader:
                self.train_data_iterator.append(iter(train_data_loader))
                self.epoch_counter.append(0)

        # Logging INFO
        path_str = [[node.name for node in path] for path in self._paths]
        children_str = [node.name for node in self._children]
        parents_str = [node.name for node in self._parents]
        if self.is_rank_zero():
            self.logger.info("*** Problem Information ***")
            self.logger.info(f"Name: {self._name}")
            self.logger.info(f"Uppers: {parents_str}")
            self.logger.info(f"Lowers: {children_str}")
            self.logger.info(f"Paths: {path_str}\n")

    def patch_everything(self):
        """
        We patch module, optimizer, data loader, and lr scheduler for device placement,
        distributed training, zero optimizer, fsdp, etc.
        """
        self.patch_module()
        self.patch_optimizer()
        if self.scheduler is not None:
            self.patch_scheduler()
        if self.train_data_loader is not None:
            self.train_data_loader = [
                self.patch_data_loader(data_loader)
                for data_loader in self.train_data_loader
            ]

    def patch_module(self):
        """
        Patch module given the systems configuration (e.g., DDP, FSDP)
        """
        self.module.to(self.device)
        if self._strategy in ["distributed", "zero"]:
            self.synchronize_params(self.parameters())
            self.module = torch.nn.parallel.DistributedDataParallel(
                module=self.module,
                gradient_as_bucket_view=True,
            )
        elif self._strategy == "fsdp":
            if self.is_rank_zero():
                self.logger.warning("FSDP requires PyTorch version >= 1.12")
            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

            self.synchronize_params(self.parameters())
            self.module = FSDP(self.module, device_id=self.device)
        elif self._strategy == "accelerate":
            self.module = self.accelerator.prepare(self.module)

    def patch_optimizer(self):
        """
        Patch optimizer given the systems configuration (e.g., DDP, FSDP)
        """
        params = self.trainable_parameters()
        if self.is_implemented("param_groups") and self._strategy != "fsdp":
            params = self.param_groups()
        is_zero = True if self._strategy == "zero" else False
        if self._strategy == "accelerate":
            self.optimizer = self.accelerator.prepare(self.optimizer)
        else:
            self.optimizer = patch_optimizer(self.optimizer, params, is_zero)

    def patch_scheduler(self):
        """
        Patch scheduler given the systems configuration (e.g., DDP, FSDP)
        """
        self.scheduler = patch_scheduler(self.scheduler, self.optimizer)
        if self._strategy == "accelerate":
            self.scheduler = self.accelerator.prepare(self.scheduler)

    def patch_data_loader(self, loader):
        """
        Patch data loader given the systems configuration (e.g., DDP, FSDP)
        """
        if self._strategy in ["distributed", "zero", "fsdp"]:
            patched_loader = get_distributed_data_loader(
                loader, world_size=self._world_size, rank=self._rank
            )
        elif self._strategy == "accelerate":
            patched_loader = self.accelerator.prepare(loader)
        else:
            patched_loader = loader

        return patched_loader

    def set_module(self, module):
        """
        Set new module for the current Problem class.
        """
        self.module = module
        self.patch_module()

    def set_optimizer(self, optimizer):
        """
        Set new optimizer for the current Problem class.
        """
        self.optimizer = optimizer
        self.patch_optimizer()

    def set_scheduler(self, scheduler):
        """
        Set new scheduler for the current Problem class.
        """
        self.scheduler = scheduler
        self.patch_scheduler()

    def set_train_data_loader(self, loader, idx=0):
        """
        Set new data loader for the current Problem class.
        """
        self.train_data_loader[idx] = self.patch_data_loader(loader)

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def forward(self, *args, **kwargs):
        """
        Users define how forward (or call) function is defined for the problem here.
        """
        return self.module(*args, **kwargs)

    @abc.abstractmethod
    def training_step(self, batch):
        """
        Users define the loss function of the problem here.
        """
        raise NotImplementedError

    def training_step_exec(self, batch):
        if self._is_default_fp16():
            with torch.cuda.amp.autocast():
                return self.training_step(batch)
        else:
            return self.training_step(batch)

    def one_step_descent(self, batch=None):
        # load data
        if batch is None:
            self.cur_batch = self.get_batch()
            batch = self.cur_batch

        # calculate loss
        loss, loss_dict = self.get_loss(batch)

        # calculate gradient (a.k.a backward)
        self.backward(
            loss=loss,
            params=self.trainable_parameters(),
            paths=self._paths,
            create_graph=not self._first_order,
            retain_graph=self._retain_graph,
            allow_unused=self._allow_unused,
        )
        if self.is_implemented("grad_callback"):
            self.grad_callback()

        # calculate parameter update
        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