# Path: ehr_prediction_modeling/types.py
# class ForwardFunctionReturns(
#     collections.namedtuple(
#         "ForwardFunctionReturns",
#         ["model_output", "hidden_states", "inputs", "activations"],
#     )):
# class EmbeddingType(object):
# class EmbeddingEncoderType(object):
# class ActivationFunction(object):
# class RegularizationType(object):
# class TaskType(object):
# class BestModelMetric(object):
# class TaskNames(object):
# class TaskTypes(object):
# class TaskLossType(object):
# class RNNCellType(object):
# class FeatureTypes(object):
# class Optimizers(object):
# class LearningRateScheduling(object):
# class EmbeddingCombinationMethod(object):
# class SNRBlockConnType(object):
# class SubNettoSubNetConnType(object):
# class SNRInputCombinationType(object):
# class LossCombinationType(object):
# class ModelTypes(object):
# class TaskLayerTypes(object):
# class ModelMode(object):
#   LOOKUP = "Lookup"
#   DEEP_EMBEDDING = "DeepEmbedding"
#   RESIDUAL = "residual"
#   FC = "fc"
#   SNR = "snr"
#   TANH = "tanh"
#   RELU = "relu"
#   LRELU = "lrelu"
#   SWISH = "swish"
#   ELU = "elu"
#   SLU = "selu"
#   ELISH = "elish"
#   HARD_ELISH = "hard_elish"
#   SIGMOID = "sigmoid"
#   HARD_SIGMOID = "hard_sigmoid"
#   TANH_PEN = "tanh_pen"
#   NONE = "None"
#   L1 = "L1"
#   L2 = "L2"
#   BINARY_CLASSIFICATION = "BinaryClassification"
#   REGRESSION = "Regression"
#   PRAUC = "prauc"
#   ROCAUC = "rocauc"
#   L1 = "l1"
#   ADVERSE_OUTCOME_RISK = "AdverseOutcomeRisk"
#   LAB_REGRESSION = "Labs"
#   BINARIZED_LOS = "BinarizedLengthOfStay"
#   REGRESSION_LOS = "RegressionLengthOfStay"
#   MORTALITY = "MortalityRisk"
#   READMISSION = "ReadmissionRisk"
#   ADVERSE_OUTCOME_RISK = "ao_risk"
#   LAB_REGRESSION = "labs_regression"
#   LOS = "length_of_stay"
#   MORTALITY = "mortality"
#   READMISSION = "readmission"
#   CE = "CE"
#   MULTI_CE = "MULTI_CE"
#   BRIER = "Brier"
#   L1 = "L1"
#   L2 = "L2"
#   SRU = "SRU"
#   LSTM = "LSTM"
#   UGRNN = "UGRNN"
#   PRESENCE_SEQ = "pres_s"  # Sequential presence features
#   NUMERIC_SEQ = "num_s"  # Sequential numerical features
#   CATEGORY_COUNTS_SEQ = "count_s"  # Sequential category counts features
# FEAT_TO_NAME = {
#     "pres_s": "presence",
#     "num_s": "numeric",
#     "count_s": "category_counts",
# }
#   SGD = "SGD"
#   ADAM = "Adam"
#   RMSPROP = "RMSprop"
#   FIXED = "fixed"
#   EXPONENTIAL_DECAY = "exponential_decay"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum_all"
#   SUM_BY_SUFFIX = "sum_by_suffix"
#   COMBINE_SNR_OUT = "combine_snr_out"
#   FC = "fc"
#   HIGHWAY = "highway"
#   RESIDUAL = "residual"
#   NONE = "None"
#   BOOL = "bool"
#   SCALAR_WEIGHT = "scalar_weight"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum"
#   SUM_ALL = "SUM_ALL"
#   UNCERTAINTY_WEIGHTED = "UNCERTAINTY_WEIGHTED"
#   RNN = "RNN"
#   SNRNN = "SNRNN"
#   NONE = "none"
#   MLP = "MLP"
#   SNRMLP = "SNRMLP"
#   TRAIN = "train"
#   EVAL = "eval"
#   PREDICT = "predict"
#   def is_train(cls, mode: str) -> bool:


from typing import Optional
from ehr_prediction_modeling import types
import tensorflow.compat.v1 as tf
            targets: tf.Tensor,
            loss_mask: tf.Tensor,
            loss_type: str = types.TaskLossType.CE,
            scale_pos_weight: Optional[float] = None) -> tf.Tensor:
  """Return the batch loss vector masked by loss_mask.

  Args:
    y: Output of the model used to compare against targets, has shape
      [num_unroll, batch_size, channels, num_targets].
    targets: Ground-truth labels used to compare against y, has shape
      [num_unroll, batch_size, channels, num_targets].
    loss_mask: Mask to apply to loss, has shape [num_unroll, batch_size,
      channels, num_targets].
    loss_type: Type of loss function. Must be one of types.TaskLossType.
    scale_pos_weight: Weight to scale positive examples by.
  Returns:
    A loss graph variable.
  """
  # TODO(b/158257206): Add a test to verify the loss is computed appropriatly
  # without silent broadcasting.
  if loss_type == types.TaskLossType.CE:

    def _loss_function(x, y, z):
      return tf.losses.sigmoid_cross_entropy(
          x, y, weights=scale_pos_fn(x, z, scale_pos_weight))
  elif loss_type == types.TaskLossType.MULTI_CE:

    def _loss_function(x, y, z):
      return tf.losses.softmax_cross_entropy(
          x, y, weights=scale_pos_fn(x, z, scale_pos_weight))

# Path: ehr_prediction_modeling/types.py
# class ForwardFunctionReturns(
#     collections.namedtuple(
#         "ForwardFunctionReturns",
#         ["model_output", "hidden_states", "inputs", "activations"],
#     )):
# class EmbeddingType(object):
# class EmbeddingEncoderType(object):
# class ActivationFunction(object):
# class RegularizationType(object):
# class TaskType(object):
# class BestModelMetric(object):
# class TaskNames(object):
# class TaskTypes(object):
# class TaskLossType(object):
# class RNNCellType(object):
# class FeatureTypes(object):
# class Optimizers(object):
# class LearningRateScheduling(object):
# class EmbeddingCombinationMethod(object):
# class SNRBlockConnType(object):
# class SubNettoSubNetConnType(object):
# class SNRInputCombinationType(object):
# class LossCombinationType(object):
# class ModelTypes(object):
# class TaskLayerTypes(object):
# class ModelMode(object):
#   LOOKUP = "Lookup"
#   DEEP_EMBEDDING = "DeepEmbedding"
#   RESIDUAL = "residual"
#   FC = "fc"
#   SNR = "snr"
#   TANH = "tanh"
#   RELU = "relu"
#   LRELU = "lrelu"
#   SWISH = "swish"
#   ELU = "elu"
#   SLU = "selu"
#   ELISH = "elish"
#   HARD_ELISH = "hard_elish"
#   SIGMOID = "sigmoid"
#   HARD_SIGMOID = "hard_sigmoid"
#   TANH_PEN = "tanh_pen"
#   NONE = "None"
#   L1 = "L1"
#   L2 = "L2"
#   BINARY_CLASSIFICATION = "BinaryClassification"
#   REGRESSION = "Regression"
#   PRAUC = "prauc"
#   ROCAUC = "rocauc"
#   L1 = "l1"
#   ADVERSE_OUTCOME_RISK = "AdverseOutcomeRisk"
#   LAB_REGRESSION = "Labs"
#   BINARIZED_LOS = "BinarizedLengthOfStay"
#   REGRESSION_LOS = "RegressionLengthOfStay"
#   MORTALITY = "MortalityRisk"
#   READMISSION = "ReadmissionRisk"
#   ADVERSE_OUTCOME_RISK = "ao_risk"
#   LAB_REGRESSION = "labs_regression"
#   LOS = "length_of_stay"
#   MORTALITY = "mortality"
#   READMISSION = "readmission"
#   CE = "CE"
#   MULTI_CE = "MULTI_CE"
#   BRIER = "Brier"
#   L1 = "L1"
#   L2 = "L2"
#   SRU = "SRU"
#   LSTM = "LSTM"
#   UGRNN = "UGRNN"
#   PRESENCE_SEQ = "pres_s"  # Sequential presence features
#   NUMERIC_SEQ = "num_s"  # Sequential numerical features
#   CATEGORY_COUNTS_SEQ = "count_s"  # Sequential category counts features
# FEAT_TO_NAME = {
#     "pres_s": "presence",
#     "num_s": "numeric",
#     "count_s": "category_counts",
# }
#   SGD = "SGD"
#   ADAM = "Adam"
#   RMSPROP = "RMSprop"
#   FIXED = "fixed"
#   EXPONENTIAL_DECAY = "exponential_decay"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum_all"
#   SUM_BY_SUFFIX = "sum_by_suffix"
#   COMBINE_SNR_OUT = "combine_snr_out"
#   FC = "fc"
#   HIGHWAY = "highway"
#   RESIDUAL = "residual"
#   NONE = "None"
#   BOOL = "bool"
#   SCALAR_WEIGHT = "scalar_weight"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum"
#   SUM_ALL = "SUM_ALL"
#   UNCERTAINTY_WEIGHTED = "UNCERTAINTY_WEIGHTED"
#   RNN = "RNN"
#   SNRNN = "SNRNN"
#   NONE = "none"
#   MLP = "MLP"
#   SNRMLP = "SNRMLP"
#   TRAIN = "train"
#   EVAL = "eval"
#   PREDICT = "predict"
#   def is_train(cls, mode: str) -> bool:


# Path: ehr_prediction_modeling/utils/loss_utils.py
from typing import Optional
from ehr_prediction_modeling import types
import tensorflow.compat.v1 as tf

            targets: tf.Tensor,
            loss_mask: tf.Tensor,
            loss_type: str = types.TaskLossType.CE,
            scale_pos_weight: Optional[float] = None) -> tf.Tensor:
  """Return the batch loss vector masked by loss_mask.

  Args:
    y: Output of the model used to compare against targets, has shape
      [num_unroll, batch_size, channels, num_targets].
    targets: Ground-truth labels used to compare against y, has shape
      [num_unroll, batch_size, channels, num_targets].
    loss_mask: Mask to apply to loss, has shape [num_unroll, batch_size,
      channels, num_targets].
    loss_type: Type of loss function. Must be one of types.TaskLossType.
    scale_pos_weight: Weight to scale positive examples by.
  Returns:
    A loss graph variable.
  """
  # TODO(b/158257206): Add a test to verify the loss is computed appropriatly
  # without silent broadcasting.
  if loss_type == types.TaskLossType.CE:

    def _loss_function(x, y, z):
      return tf.losses.sigmoid_cross_entropy(
          x, y, weights=scale_pos_fn(x, z, scale_pos_weight))
  elif loss_type == types.TaskLossType.MULTI_CE:

    def _loss_function(x, y, z):
      return tf.losses.softmax_cross_entropy(
          x, y, weights=scale_pos_fn(x, z, scale_pos_weight))

  elif loss_type == types.TaskLossType.BRIER: