# Path: tensorflow_lattice/python/aggregation_layer.py
# class Aggregation(keras.layers.Layer):
#   def __init__(self, model, **kwargs):
#   def call(self, x):
#   def get_config(self):
#   def from_config(cls, config, custom_objects=None):


import tensorflow as tf
from tensorflow_lattice.python import aggregation_layer
# 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.
"""Tests for Tensorflow Lattice premade."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


test_input = [
    tf.ragged.constant([[1, 2], [1, 2, 3], [3]]),
    tf.ragged.constant([[4, 5], [4, 4, 4], [6]]),
    tf.ragged.constant([[1, 6], [5, 5, 5], [9]])
]

expected_output = tf.constant([32, 40, 162])


class AggregationTest(tf.test.TestCase):

  def testAggregationLayer(self):
    # First we test our assertion that the model must be a tf.keras.Model
    with self.assertRaisesRegex(ValueError,
                                'Model must be a tf.keras.Model instance.'):
      aggregation_layer.Aggregation(None)
    # Now let's make sure our layer aggregates properly.
    inputs = [tf.keras.Input(shape=()) for _ in range(len(test_input))]
    output = tf.keras.layers.multiply(inputs)
    model = tf.keras.Model(inputs=inputs, outputs=output)

# Path: tensorflow_lattice/python/aggregation_layer.py
# class Aggregation(keras.layers.Layer):
#   def __init__(self, model, **kwargs):
#   def call(self, x):
#   def get_config(self):
#   def from_config(cls, config, custom_objects=None):


# Path: tensorflow_lattice/python/aggregation_test.py
import tensorflow as tf
from tensorflow_lattice.python import aggregation_layer

# 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.
"""Tests for Tensorflow Lattice premade."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


test_input = [
    tf.ragged.constant([[1, 2], [1, 2, 3], [3]]),
    tf.ragged.constant([[4, 5], [4, 4, 4], [6]]),
    tf.ragged.constant([[1, 6], [5, 5, 5], [9]])
]

expected_output = tf.constant([32, 40, 162])


class AggregationTest(tf.test.TestCase):

  def testAggregationLayer(self):
    # First we test our assertion that the model must be a tf.keras.Model
    with self.assertRaisesRegex(ValueError,
                                'Model must be a tf.keras.Model instance.'):
      aggregation_layer.Aggregation(None)
    # Now let's make sure our layer aggregates properly.
    inputs = [tf.keras.Input(shape=()) for _ in range(len(test_input))]
    output = tf.keras.layers.multiply(inputs)
    model = tf.keras.Model(inputs=inputs, outputs=output)

    agg_layer = aggregation_layer.Aggregation(model)