# Path: ehr_prediction_modeling/data/tf_dataset_utils.py
# def sample_segments(
#     seq_lens: tf.Tensor,
#     segment_length: tf.Tensor,
#     random_offset: bool,
# ) -> tf.Tensor:
#   def flatten(x):
# def slice_tensor(
#     tensor: tf.Tensor,
#     seq_starts: tf.Tensor,
#     segment_length: int,
# ) -> tf.Tensor:
#   def loop_body(i, ta):
# def slice_sparse_tensor_sequences(
#     tensor: tf.SparseTensor,
#     seq_starts: tf.Tensor,
#     segment_length: int,
# ) -> tf.SparseTensor:
# def slice_sparse_tensor_context(tensor: tf.SparseTensor,
#                                 seq_starts: tf.Tensor) -> tf.SparseTensor:
# def convert_to_segments(
#     examples: SeqexDictType,
#     segment_len: int,
# ) -> SeqexDictType:
# def read_seqex_dataset(tfrecords_filename: str) -> tf.data.TFRecordDataset:
# def tile_context(
#     ctx: Dict[str, tf.Tensor],
#     seq_len: int,
#     batch_padding: Optional[int] = None,
# ) -> Dict[str, tf.Tensor]:
# def add_beginning_of_sequence_mask_and_pad(
#     example: SeqexDictType,
#     batch_size: int,
#     num_unroll: int,
# ) -> tf.data.Dataset:
# def seqex_to_dict(
#     serialized_examples: bytes,
#     context_d: Dict[str, str],
#     sequence_d: Dict[str, str],
# ) -> SeqexDictType:
# def get_label_dicts(
#     task_coordinator: coordinator.Coordinator,
#     context_features: List[str],
#     sequential_features: List[str],
# ) -> Tuple[Dict[str, Union[tf.FixedLenSequenceFeature, tf.FixedLenFeature]],
# def to_time_major(tensor: tf.Tensor) -> tf.Tensor:
#   def _transpose(value, dimension):
# def sparse_tensor_to_time_major(
#     sparse_tensor: tf.SparseTensor) -> tf.SparseTensor:
# def convert_to_time_major(examples: SeqexDictType) -> SeqexDictType:
# 
# Path: ehr_prediction_modeling/tasks/coordinator.py
# class Coordinator():
#   def __init__(self, task_list: List[base_task.Task],
#                optimizer_config: configdict.ConfigDict):
#   def num_targets_list(self) -> List[int]:
#   def task_list(self) -> List[base_task.Task]:
#   def task_names(self) -> List[str]:
#   def target_names_list(self) -> List[str]:
#   def task_prediction_types(self) -> List[str]:
#   def task_eval_masks_list(self) -> List[List[str]]:
#   def task_layers(self) -> Mapping[str, Any]:
#   def get_task_layers_variables(self) -> Sequence[tf.Variable]:
#   def get_coordinator_variables(
#       self, batch: batches.TFBatch,
#       model_output: Union[tf.Tensor, None]) -> task_data.CoordinatorVariables:
#   def get_task_regularization_losses(self) -> tf.Tensor:
#   def _get_task_variables_and_loss(
#       self, batch: batches.TFBatch,
#       model_output: tf.Tensor) -> task_data.CoordinatorVariables:
#   def _get_task_variables(
#       self, batch: batches.TFBatch) -> task_data.CoordinatorVariables:
#   def get_label_dicts(
#       self
#   ) -> Tuple[Dict[str, Union[tf.FixedLenSequenceFeature, tf.FixedLenFeature]],
# 
# Path: ehr_prediction_modeling/utils/batches.py
# def batch_to_components(
#     batch: TFBatch,
#     context_features: List[str],
#     sequential_features: List[str],
# ) -> Dict[str, tf.Tensor]:
# def _batch_to_components(
#     batch: TFBatch,
#     context_features: List[str],
#     sequential_features: List[str],
# ) -> Dict[str, tf.Tensor]:
# def _extract_features(
#     batch_field: Dict[str, tf.Tensor],
#     features: List[str],
#     outp_prefix: str,
#     inp_prefix: str,
# ) -> Dict[str, tf.Tensor]:
# def flatten_batch(val: tf.Tensor) -> tf.Tensor:
# def get_unflatten_batch_fn(
#     batch_shape: tf.Tensor) -> Callable[[tf.Tensor], tf.Tensor]:
#   def unflatten_batch(val):
# def sparse_fill_empty_rows(
#     val: Union[tf.Tensor,
#                tf.SparseTensor]) -> Union[tf.Tensor, tf.SparseTensor]:
# 
# Path: ehr_prediction_modeling/configdict.py
# class ConfigDict(dict):
#   def __init__(self, config=None, **kwargs):  # pylint: disable=super-init-not-called
#   def __setattr__(self, key, value):
#   def __getitem__(self, key):
#   def update(self, config=None, **f):
#   def get(self, key, default=None):


import os
import tensorflow.compat.v1 as tf
from typing import List, Optional
from absl import logging
from ehr_prediction_modeling.data import tf_dataset_utils
from ehr_prediction_modeling.tasks import coordinator
from ehr_prediction_modeling.utils import batches
from ehr_prediction_modeling import configdict
# coding=utf-8
# Copyright 2021 Google Health Research.
#
# 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.

# Lint as: python3
"""TensorFlow utility functions to extract Tensor and batch from tf protos."""





class BatchGenerator(object):
  """Tool for working with a TF dataset."""

  def __init__(
      self,
      config: configdict.ConfigDict,
      is_training: bool,

# Path: ehr_prediction_modeling/data/tf_dataset_utils.py
# def sample_segments(
#     seq_lens: tf.Tensor,
#     segment_length: tf.Tensor,
#     random_offset: bool,
# ) -> tf.Tensor:
#   def flatten(x):
# def slice_tensor(
#     tensor: tf.Tensor,
#     seq_starts: tf.Tensor,
#     segment_length: int,
# ) -> tf.Tensor:
#   def loop_body(i, ta):
# def slice_sparse_tensor_sequences(
#     tensor: tf.SparseTensor,
#     seq_starts: tf.Tensor,
#     segment_length: int,
# ) -> tf.SparseTensor:
# def slice_sparse_tensor_context(tensor: tf.SparseTensor,
#                                 seq_starts: tf.Tensor) -> tf.SparseTensor:
# def convert_to_segments(
#     examples: SeqexDictType,
#     segment_len: int,
# ) -> SeqexDictType:
# def read_seqex_dataset(tfrecords_filename: str) -> tf.data.TFRecordDataset:
# def tile_context(
#     ctx: Dict[str, tf.Tensor],
#     seq_len: int,
#     batch_padding: Optional[int] = None,
# ) -> Dict[str, tf.Tensor]:
# def add_beginning_of_sequence_mask_and_pad(
#     example: SeqexDictType,
#     batch_size: int,
#     num_unroll: int,
# ) -> tf.data.Dataset:
# def seqex_to_dict(
#     serialized_examples: bytes,
#     context_d: Dict[str, str],
#     sequence_d: Dict[str, str],
# ) -> SeqexDictType:
# def get_label_dicts(
#     task_coordinator: coordinator.Coordinator,
#     context_features: List[str],
#     sequential_features: List[str],
# ) -> Tuple[Dict[str, Union[tf.FixedLenSequenceFeature, tf.FixedLenFeature]],
# def to_time_major(tensor: tf.Tensor) -> tf.Tensor:
#   def _transpose(value, dimension):
# def sparse_tensor_to_time_major(
#     sparse_tensor: tf.SparseTensor) -> tf.SparseTensor:
# def convert_to_time_major(examples: SeqexDictType) -> SeqexDictType:
# 
# Path: ehr_prediction_modeling/tasks/coordinator.py
# class Coordinator():
#   def __init__(self, task_list: List[base_task.Task],
#                optimizer_config: configdict.ConfigDict):
#   def num_targets_list(self) -> List[int]:
#   def task_list(self) -> List[base_task.Task]:
#   def task_names(self) -> List[str]:
#   def target_names_list(self) -> List[str]:
#   def task_prediction_types(self) -> List[str]:
#   def task_eval_masks_list(self) -> List[List[str]]:
#   def task_layers(self) -> Mapping[str, Any]:
#   def get_task_layers_variables(self) -> Sequence[tf.Variable]:
#   def get_coordinator_variables(
#       self, batch: batches.TFBatch,
#       model_output: Union[tf.Tensor, None]) -> task_data.CoordinatorVariables:
#   def get_task_regularization_losses(self) -> tf.Tensor:
#   def _get_task_variables_and_loss(
#       self, batch: batches.TFBatch,
#       model_output: tf.Tensor) -> task_data.CoordinatorVariables:
#   def _get_task_variables(
#       self, batch: batches.TFBatch) -> task_data.CoordinatorVariables:
#   def get_label_dicts(
#       self
#   ) -> Tuple[Dict[str, Union[tf.FixedLenSequenceFeature, tf.FixedLenFeature]],
# 
# Path: ehr_prediction_modeling/utils/batches.py
# def batch_to_components(
#     batch: TFBatch,
#     context_features: List[str],
#     sequential_features: List[str],
# ) -> Dict[str, tf.Tensor]:
# def _batch_to_components(
#     batch: TFBatch,
#     context_features: List[str],
#     sequential_features: List[str],
# ) -> Dict[str, tf.Tensor]:
# def _extract_features(
#     batch_field: Dict[str, tf.Tensor],
#     features: List[str],
#     outp_prefix: str,
#     inp_prefix: str,
# ) -> Dict[str, tf.Tensor]:
# def flatten_batch(val: tf.Tensor) -> tf.Tensor:
# def get_unflatten_batch_fn(
#     batch_shape: tf.Tensor) -> Callable[[tf.Tensor], tf.Tensor]:
#   def unflatten_batch(val):
# def sparse_fill_empty_rows(
#     val: Union[tf.Tensor,
#                tf.SparseTensor]) -> Union[tf.Tensor, tf.SparseTensor]:
# 
# Path: ehr_prediction_modeling/configdict.py
# class ConfigDict(dict):
#   def __init__(self, config=None, **kwargs):  # pylint: disable=super-init-not-called
#   def __setattr__(self, key, value):
#   def __getitem__(self, key):
#   def update(self, config=None, **f):
#   def get(self, key, default=None):


# Path: ehr_prediction_modeling/data/tf_dataset.py
import os
import tensorflow.compat.v1 as tf
from typing import List, Optional
from absl import logging
from ehr_prediction_modeling.data import tf_dataset_utils
from ehr_prediction_modeling.tasks import coordinator
from ehr_prediction_modeling.utils import batches
from ehr_prediction_modeling import configdict

# coding=utf-8
# Copyright 2021 Google Health Research.
#
# 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.

# Lint as: python3
"""TensorFlow utility functions to extract Tensor and batch from tf protos."""





class BatchGenerator(object):
  """Tool for working with a TF dataset."""

  def __init__(
      self,
      config: configdict.ConfigDict,
      is_training: bool,

      task_coordinator: Optional[coordinator.Coordinator],