# Path: tensorflow_lattice/python/cdf_layer.py
# class CDF(tf.keras.layers.Layer):
#   def __init__(self,
#                num_keypoints,
#                units=1,
#                activation="relu6",
#                reduction="mean",
#                input_scaling_init=None,
#                input_scaling_type="fixed",
#                input_scaling_monotonicity="increasing",
#                sparsity_factor=1,
#                kernel_initializer="random_uniform",
#                **kwargs):
#   def build(self, input_shape):
#   def call(self, inputs):
#   def get_config(self):
# def create_kernel_initializer(kernel_initializer_id):
# 
# Path: tensorflow_lattice/python/test_utils.py
# class TimeTracker(object):
#   def __init__(self, list_to_append, num_steps=1):
#   def __enter__(self):
#   def __exit__(self, unuesd_type, unuesd_value, unuesd_traceback):
# def run_training_loop(config,
#                       training_data,
#                       keras_model,
#                       plot_path=None,
#                       input_dtype=np.float32,
#                       label_dtype=np.float32):
# def assert_sonnet_equivalent_to_keras(
#     test, sonnet_module_ctor, keras_layer_ctor,
#     training_inputs, training_labels,
#     epsilon=1e-4):
# def two_dim_mesh_grid(num_points, x_min, y_min, x_max, y_max):
# def sample_uniformly(num_points, lower_bounds, upper_bounds):
# def get_hypercube_interpolation_fn(coefficients):
#   def hypercube_interpolation_fn(x):
# def get_linear_lattice_interpolation_fn(lattice_sizes, monotonicities,
#                                         output_min, output_max):
#   def linear_interpolation_fn(x):


import math
import numpy as np
import tensorflow as tf
from absl import logging
from absl.testing import parameterized
from tensorflow_lattice.python import cdf_layer
from tensorflow_lattice.python import test_utils

    Returns:
      Tuple `(training_inputs, training_labels, raw_training_inputs)` where
        `training_inputs` and `training_labels` are data for training and
        `raw_training_inputs` are representation of training_inputs for
        visualization.
    """
    raw_training_inputs = config["x_generator"](
        num_points=config["num_training_records"],
        input_dims=config["input_dims"])

    if isinstance(raw_training_inputs, tuple):
      # This means that raw inputs are 2-d mesh grid. Convert them into list of
      # 2-d points.
      training_inputs = list(np.dstack(raw_training_inputs).reshape((-1, 2)))
    else:
      training_inputs = raw_training_inputs

    training_labels = [config["y_function"](x) for x in training_inputs]
    return training_inputs, training_labels, raw_training_inputs

  def _TrainModel(self, config, plot_path=None):
    logging.info("Testing config:")
    logging.info(config)
    config = self._SetDefaults(config)
    self._ResetAllBackends()

    training_inputs, training_labels, raw_training_inputs = (
        self._GetTrainingInputsAndLabels(config))


# Path: tensorflow_lattice/python/cdf_layer.py
# class CDF(tf.keras.layers.Layer):
#   def __init__(self,
#                num_keypoints,
#                units=1,
#                activation="relu6",
#                reduction="mean",
#                input_scaling_init=None,
#                input_scaling_type="fixed",
#                input_scaling_monotonicity="increasing",
#                sparsity_factor=1,
#                kernel_initializer="random_uniform",
#                **kwargs):
#   def build(self, input_shape):
#   def call(self, inputs):
#   def get_config(self):
# def create_kernel_initializer(kernel_initializer_id):
# 
# Path: tensorflow_lattice/python/test_utils.py
# class TimeTracker(object):
#   def __init__(self, list_to_append, num_steps=1):
#   def __enter__(self):
#   def __exit__(self, unuesd_type, unuesd_value, unuesd_traceback):
# def run_training_loop(config,
#                       training_data,
#                       keras_model,
#                       plot_path=None,
#                       input_dtype=np.float32,
#                       label_dtype=np.float32):
# def assert_sonnet_equivalent_to_keras(
#     test, sonnet_module_ctor, keras_layer_ctor,
#     training_inputs, training_labels,
#     epsilon=1e-4):
# def two_dim_mesh_grid(num_points, x_min, y_min, x_max, y_max):
# def sample_uniformly(num_points, lower_bounds, upper_bounds):
# def get_hypercube_interpolation_fn(coefficients):
#   def hypercube_interpolation_fn(x):
# def get_linear_lattice_interpolation_fn(lattice_sizes, monotonicities,
#                                         output_min, output_max):
#   def linear_interpolation_fn(x):


# Path: tensorflow_lattice/python/cdf_test.py
import math
import numpy as np
import tensorflow as tf
from absl import logging
from absl.testing import parameterized
from tensorflow_lattice.python import cdf_layer
from tensorflow_lattice.python import test_utils


    Returns:
      Tuple `(training_inputs, training_labels, raw_training_inputs)` where
        `training_inputs` and `training_labels` are data for training and
        `raw_training_inputs` are representation of training_inputs for
        visualization.
    """
    raw_training_inputs = config["x_generator"](
        num_points=config["num_training_records"],
        input_dims=config["input_dims"])

    if isinstance(raw_training_inputs, tuple):
      # This means that raw inputs are 2-d mesh grid. Convert them into list of
      # 2-d points.
      training_inputs = list(np.dstack(raw_training_inputs).reshape((-1, 2)))
    else:
      training_inputs = raw_training_inputs

    training_labels = [config["y_function"](x) for x in training_inputs]
    return training_inputs, training_labels, raw_training_inputs

  def _TrainModel(self, config, plot_path=None):
    logging.info("Testing config:")
    logging.info(config)
    config = self._SetDefaults(config)
    self._ResetAllBackends()

    training_inputs, training_labels, raw_training_inputs = (
        self._GetTrainingInputsAndLabels(config))


    keras_layer = cdf_layer.CDF(