# Path: ehr_prediction_modeling/embeddings/input_embedding_base.py
# class InputEmbeddingBase(metaclass=abc.ABCMeta):
#   def __init__(self,
#                encoder_config: "configdict.ConfigDict",
#                emb_dim_dict: Dict[str, int],
#                has_embedding_loss: Optional[bool] = False):
#   def get_encoders(self):
#   def _initialize_weights(self):
#   def get_embedding_loss(self, *unused_args, **unused_kwargs
#                          ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
#   def _embed_feature(
#       self,
#       feat_type: str,
#       input_data: Tuple[tf.SparseTensor, tf.SparseTensor],
#       context: Optional[Dict[str, Any]] = None
#       ) -> tf.Tensor:
#   def provide_embeddings_to_forward_fn(
#       self,
#       data: Dict[str, tf.SparseTensor],
#       feature_types: List[str]
#       ) -> tf.Tensor:
# 
# Path: ehr_prediction_modeling/models/nets/sparse_lookup.py
# VALID_COMBINERS = ["sum", "sqrtn", "mean"]
# def combiner_correction(sparse_combiner: str, embed: tf.Tensor) -> tf.Tensor:
#   def __init__(self,
#                ndim_emb: int,
#                ndim_input: int,
#                n_act: Optional[int] = None,
#                activate_final: bool = False,
#                act_fn: Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh,
#                sparse_lookup_dropout_prob: float = 0.,
#                dropout_is_training: bool = True,
#                sparse_combine: str = "sum",
#                name: Optional[str] = None,
#                identity_lookup: bool = False):
#   def _build(
#       self,
#       input_data: Tuple[tf.SparseTensor, tf.SparseTensor],
#       context: Optional[Dict[str, Union[tf.Tensor, bool,
#                                         tf.SparseTensor]]] = None
#   ) -> tf.Tensor:
# class SparseLookupEncoder(snt.AbstractModule):


import typing
import tensorflow.compat.v1 as tf
from typing import Dict
from ehr_prediction_modeling.embeddings import input_embedding_base
from ehr_prediction_modeling.models.nets import sparse_lookup
  from ehr_prediction_modeling import configdict
# 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.

# Lint as: python3
"""Lookup Embedding module."""




if typing.TYPE_CHECKING:


class BasicEmbeddingLookup(input_embedding_base.InputEmbeddingBase):
  """The class that uses the base embedding lookup."""

  def __init__(self, encoder_config: "configdict.ConfigDict",
               emb_dim_dict: Dict[str, int]):
    """Initialize the embedding base object."""
    super().__init__(encoder_config, emb_dim_dict)

    with tf.variable_scope(None, default_name=self._name) as variable_scope:
      self.variable_scope = variable_scope
      # Initialize all the weights.
      self._initialize_weights()

  def _initialize_weights(self):
    """Initialize the weights for embedding computations."""
    for feat_type in self._all_features:

# Path: ehr_prediction_modeling/embeddings/input_embedding_base.py
# class InputEmbeddingBase(metaclass=abc.ABCMeta):
#   def __init__(self,
#                encoder_config: "configdict.ConfigDict",
#                emb_dim_dict: Dict[str, int],
#                has_embedding_loss: Optional[bool] = False):
#   def get_encoders(self):
#   def _initialize_weights(self):
#   def get_embedding_loss(self, *unused_args, **unused_kwargs
#                          ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
#   def _embed_feature(
#       self,
#       feat_type: str,
#       input_data: Tuple[tf.SparseTensor, tf.SparseTensor],
#       context: Optional[Dict[str, Any]] = None
#       ) -> tf.Tensor:
#   def provide_embeddings_to_forward_fn(
#       self,
#       data: Dict[str, tf.SparseTensor],
#       feature_types: List[str]
#       ) -> tf.Tensor:
# 
# Path: ehr_prediction_modeling/models/nets/sparse_lookup.py
# VALID_COMBINERS = ["sum", "sqrtn", "mean"]
# def combiner_correction(sparse_combiner: str, embed: tf.Tensor) -> tf.Tensor:
#   def __init__(self,
#                ndim_emb: int,
#                ndim_input: int,
#                n_act: Optional[int] = None,
#                activate_final: bool = False,
#                act_fn: Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh,
#                sparse_lookup_dropout_prob: float = 0.,
#                dropout_is_training: bool = True,
#                sparse_combine: str = "sum",
#                name: Optional[str] = None,
#                identity_lookup: bool = False):
#   def _build(
#       self,
#       input_data: Tuple[tf.SparseTensor, tf.SparseTensor],
#       context: Optional[Dict[str, Union[tf.Tensor, bool,
#                                         tf.SparseTensor]]] = None
#   ) -> tf.Tensor:
# class SparseLookupEncoder(snt.AbstractModule):


# Path: ehr_prediction_modeling/embeddings/basic_embedding_lookup.py
import typing
import tensorflow.compat.v1 as tf
from typing import Dict
from ehr_prediction_modeling.embeddings import input_embedding_base
from ehr_prediction_modeling.models.nets import sparse_lookup
  from ehr_prediction_modeling import configdict

# 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.

# Lint as: python3
"""Lookup Embedding module."""




if typing.TYPE_CHECKING:


class BasicEmbeddingLookup(input_embedding_base.InputEmbeddingBase):
  """The class that uses the base embedding lookup."""

  def __init__(self, encoder_config: "configdict.ConfigDict",
               emb_dim_dict: Dict[str, int]):
    """Initialize the embedding base object."""
    super().__init__(encoder_config, emb_dim_dict)

    with tf.variable_scope(None, default_name=self._name) as variable_scope:
      self.variable_scope = variable_scope
      # Initialize all the weights.
      self._initialize_weights()

  def _initialize_weights(self):
    """Initialize the weights for embedding computations."""
    for feat_type in self._all_features:

      self._encoders[feat_type] = sparse_lookup.SparseLookupEncoder(