# 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):


import functools
import flax.linen as nn
import jax.numpy as jnp
from typing import List
from init2winit.model_lib import base_model
from jax import nn as jnn
from ml_collections.config_dict import config_dict
        scaled_variance_scaling_init,
        [self.total_vocab_sizes, self.embed_dim])

    idx_lookup = jnp.reshape(idx_lookup, [-1])
    embed_features = embedding_table[idx_lookup]
    embed_features = jnp.reshape(
        embed_features, [batch_size, -1, self.embed_dim])
    feature_stack = jnp.concatenate([feature_stack, embed_features], axis=1)
    dot_interact_output = dot_interact(
        concat_features=feature_stack, keep_diags=self.keep_diags)
    top_mlp_input = jnp.concatenate([bot_mlp_output, dot_interact_output],
                                    axis=-1)
    mlp_input_dim = top_mlp_input.shape[1]
    mlp_top_dims = self.mlp_top_dims
    num_layers_top = len(mlp_top_dims)
    for layer_idx, fan_out in enumerate(mlp_top_dims):
      fan_in = mlp_input_dim if layer_idx == 0 else mlp_top_dims[layer_idx - 1]
      top_mlp_input = nn.Dense(
          fan_out,
          kernel_init=jnn.initializers.normal(
              stddev=jnp.sqrt(2.0 / (fan_in + fan_out))),
          bias_init=jnn.initializers.normal(
              stddev=jnp.sqrt(1.0 / mlp_top_dims[layer_idx])))(
                  top_mlp_input)
      if layer_idx < (num_layers_top - 1):
        top_mlp_input = nn.relu(top_mlp_input)
    logits = top_mlp_input
    return logits



# 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/dlrm.py
import functools
import flax.linen as nn
import jax.numpy as jnp
from typing import List
from init2winit.model_lib import base_model
from jax import nn as jnn
from ml_collections.config_dict import config_dict

        scaled_variance_scaling_init,
        [self.total_vocab_sizes, self.embed_dim])

    idx_lookup = jnp.reshape(idx_lookup, [-1])
    embed_features = embedding_table[idx_lookup]
    embed_features = jnp.reshape(
        embed_features, [batch_size, -1, self.embed_dim])
    feature_stack = jnp.concatenate([feature_stack, embed_features], axis=1)
    dot_interact_output = dot_interact(
        concat_features=feature_stack, keep_diags=self.keep_diags)
    top_mlp_input = jnp.concatenate([bot_mlp_output, dot_interact_output],
                                    axis=-1)
    mlp_input_dim = top_mlp_input.shape[1]
    mlp_top_dims = self.mlp_top_dims
    num_layers_top = len(mlp_top_dims)
    for layer_idx, fan_out in enumerate(mlp_top_dims):
      fan_in = mlp_input_dim if layer_idx == 0 else mlp_top_dims[layer_idx - 1]
      top_mlp_input = nn.Dense(
          fan_out,
          kernel_init=jnn.initializers.normal(
              stddev=jnp.sqrt(2.0 / (fan_in + fan_out))),
          bias_init=jnn.initializers.normal(
              stddev=jnp.sqrt(1.0 / mlp_top_dims[layer_idx])))(
                  top_mlp_input)
      if layer_idx < (num_layers_top - 1):
        top_mlp_input = nn.relu(top_mlp_input)
    logits = top_mlp_input
    return logits



class DLRMModel(base_model.BaseModel):