# 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
      (4, 20, 10, "relu6", "mean", "fixed", 27),
      (4, 20, 10, "relu6", "mean", "learned_shared", 32),
      (4, 20, 10, "relu6", "mean", "learned_per_input", 32),
      (4, 20, 10, "relu6", "geometric_mean", "fixed", 33),
      (4, 20, 10, "relu6", "geometric_mean", "learned_shared", 38),
      (4, 20, 10, "relu6", "geometric_mean", "learned_per_input", 38),
      (2, 10, 5, "sigmoid", "mean", "fixed", 25),
      (2, 10, 5, "sigmoid", "mean", "learned_shared", 30),
      (2, 10, 5, "sigmoid", "mean", "learned_per_input", 30),
      (2, 10, 5, "sigmoid", "geometric_mean", "fixed", 31),
      (2, 10, 5, "sigmoid", "geometric_mean", "learned_shared", 36),
      (2, 10, 5, "sigmoid", "geometric_mean", "learned_per_input", 36),
      (4, 20, 10, "sigmoid", "mean", "fixed", 25),
      (4, 20, 10, "sigmoid", "mean", "learned_shared", 30),
      (4, 20, 10, "sigmoid", "mean", "learned_per_input", 30),
      (4, 20, 10, "sigmoid", "geometric_mean", "fixed", 31),
      (4, 20, 10, "sigmoid", "geometric_mean", "learned_shared", 36),
      (4, 20, 10, "sigmoid", "geometric_mean", "learned_per_input", 36),
  )
  def testGraphSize(self, input_dims, num_keypoints, units, activation,
                    reduction, input_scaling_type, expected_graph_size):
    # If this test failed then you modified core lattice interpolation logic in
    # a way which increases number of ops in the graph. Or maybe Keras team
    # changed something under the hood. Please ensure that this increase is
    # unavoidable and try to minimize it.
    if self.disable_all:
      return
    tf.compat.v1.disable_eager_execution()
    tf.compat.v1.reset_default_graph()


# 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

      (4, 20, 10, "relu6", "mean", "fixed", 27),
      (4, 20, 10, "relu6", "mean", "learned_shared", 32),
      (4, 20, 10, "relu6", "mean", "learned_per_input", 32),
      (4, 20, 10, "relu6", "geometric_mean", "fixed", 33),
      (4, 20, 10, "relu6", "geometric_mean", "learned_shared", 38),
      (4, 20, 10, "relu6", "geometric_mean", "learned_per_input", 38),
      (2, 10, 5, "sigmoid", "mean", "fixed", 25),
      (2, 10, 5, "sigmoid", "mean", "learned_shared", 30),
      (2, 10, 5, "sigmoid", "mean", "learned_per_input", 30),
      (2, 10, 5, "sigmoid", "geometric_mean", "fixed", 31),
      (2, 10, 5, "sigmoid", "geometric_mean", "learned_shared", 36),
      (2, 10, 5, "sigmoid", "geometric_mean", "learned_per_input", 36),
      (4, 20, 10, "sigmoid", "mean", "fixed", 25),
      (4, 20, 10, "sigmoid", "mean", "learned_shared", 30),
      (4, 20, 10, "sigmoid", "mean", "learned_per_input", 30),
      (4, 20, 10, "sigmoid", "geometric_mean", "fixed", 31),
      (4, 20, 10, "sigmoid", "geometric_mean", "learned_shared", 36),
      (4, 20, 10, "sigmoid", "geometric_mean", "learned_per_input", 36),
  )
  def testGraphSize(self, input_dims, num_keypoints, units, activation,
                    reduction, input_scaling_type, expected_graph_size):
    # If this test failed then you modified core lattice interpolation logic in
    # a way which increases number of ops in the graph. Or maybe Keras team
    # changed something under the hood. Please ensure that this increase is
    # unavoidable and try to minimize it.
    if self.disable_all:
      return
    tf.compat.v1.disable_eager_execution()
    tf.compat.v1.reset_default_graph()


    layer = cdf_layer.CDF(