# Path: init2winit/model_lib/base_model.py
# def _evaluate_batch(flax_module, params, batch_stats, batch, metrics_bundle,
#                     apply_one_hot_in_loss):
# def _predict_batch(flax_module,
#                    params,
#                    batch_stats,
#                    batch,
#                    output_activation_fn=None):
#   def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
#   def evaluate_batch(self, params, batch_stats, batch):
#   def predict_batch(self,
#                     params,
#                     batch_stats,
#                     batch,
#                     apply_output_activation_fn=False):
#   def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
#   def get_fake_batch(self, hps):
#   def build_flax_module(self):
# class BaseModel(object):


from typing import Optional, Sequence, Union
from flax import linen as nn
from init2winit.model_lib import base_model
from ml_collections.config_dict import config_dict
import jax.numpy as jnp
import numpy as np
      x = out['head_input'] = x[:, 0]
    elif self.pool_type == 'tok':
      x = out['head_input'] = x[:, 0]
      encoded = encoded[:, 1:]
    else:
      raise ValueError(f'Unknown pool type: "{self.pool_type}"')

    x_2d = jnp.reshape(encoded, [n, h, w, -1])

    if self.rep_size:
      rep_size = self.width if self.rep_size is True else self.rep_size  # pylint: disable=g-bool-id-comparison
      hid = nn.Dense(rep_size, name='pre_logits')
      # NOTE: In the past we did not include tanh in pre_logits.
      # For few-shot, it should not matter much, as it whitens anyways.
      x_2d = nn.tanh(hid(x_2d))
      x = nn.tanh(hid(x))

    out['pre_logits_2d'] = x_2d
    out['pre_logits'] = x

    if self.num_classes:
      kw = {'kernel_init': nn.initializers.zeros} if self.head_zeroinit else {}
      head = nn.Dense(self.num_classes, name='head', **kw)
      x_2d = out['logits_2d'] = head(x_2d)
      x = out['logits'] = head(x)

    # TODO(dsuo): this used to be `return x, out`. Do we need out?
    return x



# Path: init2winit/model_lib/base_model.py
# def _evaluate_batch(flax_module, params, batch_stats, batch, metrics_bundle,
#                     apply_one_hot_in_loss):
# def _predict_batch(flax_module,
#                    params,
#                    batch_stats,
#                    batch,
#                    output_activation_fn=None):
#   def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
#   def evaluate_batch(self, params, batch_stats, batch):
#   def predict_batch(self,
#                     params,
#                     batch_stats,
#                     batch,
#                     apply_output_activation_fn=False):
#   def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
#   def get_fake_batch(self, hps):
#   def build_flax_module(self):
# class BaseModel(object):


# Path: init2winit/model_lib/vit.py
from typing import Optional, Sequence, Union
from flax import linen as nn
from init2winit.model_lib import base_model
from ml_collections.config_dict import config_dict
import jax.numpy as jnp
import numpy as np

      x = out['head_input'] = x[:, 0]
    elif self.pool_type == 'tok':
      x = out['head_input'] = x[:, 0]
      encoded = encoded[:, 1:]
    else:
      raise ValueError(f'Unknown pool type: "{self.pool_type}"')

    x_2d = jnp.reshape(encoded, [n, h, w, -1])

    if self.rep_size:
      rep_size = self.width if self.rep_size is True else self.rep_size  # pylint: disable=g-bool-id-comparison
      hid = nn.Dense(rep_size, name='pre_logits')
      # NOTE: In the past we did not include tanh in pre_logits.
      # For few-shot, it should not matter much, as it whitens anyways.
      x_2d = nn.tanh(hid(x_2d))
      x = nn.tanh(hid(x))

    out['pre_logits_2d'] = x_2d
    out['pre_logits'] = x

    if self.num_classes:
      kw = {'kernel_init': nn.initializers.zeros} if self.head_zeroinit else {}
      head = nn.Dense(self.num_classes, name='head', **kw)
      x_2d = out['logits_2d'] = head(x_2d)
      x = out['logits'] = head(x)

    # TODO(dsuo): this used to be `return x, out`. Do we need out?
    return x



class ViTModel(base_model.BaseModel):