# Copyright 2017 The Sonnet Authors. All Rights Reserved.
#
# 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.
# ============================================================================

"""Tests for `sonnet.python.modules.conv`."""

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

import itertools
import random

# Dependency imports
from absl.testing import parameterized
import numpy as np
import sonnet as snt
from sonnet.python.modules import conv
import tensorflow.compat.v1 as tf
from tensorflow.contrib import layers as contrib_layers

from tensorflow.python.ops import variables  # pylint: disable=g-direct-tensorflow-import


def create_constant_initializers(w, b, use_bias):
  if use_bias:
    return {"w": tf.constant_initializer(w),
            "b": tf.constant_initializer(b)}
  else:
    return {"w": tf.constant_initializer(w)}


def create_separable_constant_initializers(w_dw, w_pw, b, use_bias):
  if use_bias:
    return {"w_dw": tf.constant_initializer(w_dw),
            "w_pw": tf.constant_initializer(w_pw),
            "b": tf.constant_initializer(b)}
  else:
    return {"w_dw": tf.constant_initializer(w_dw),
            "w_pw": tf.constant_initializer(w_pw)}


def create_regularizers(use_bias, regularizer):
  if use_bias:
    return {"w": regularizer, "b": regularizer}
  else:
    return {"w": regularizer}


def create_separable_regularizers(use_bias, regularizer):
  if use_bias:
    return {"w_dw": regularizer, "w_pw": regularizer, "b": regularizer}
  else:
    return {"w_dw": regularizer, "w_pw": regularizer}


class FillListTest(tf.test.TestCase):

  def test(self):
    """Tests the _fill_list private function in snt.conv."""
    x = random.randint(1, 10)

    self.assertEqual(conv._fill_shape(x, 1), (x,))
    self.assertEqual(conv._fill_shape(x, 2), (x, x))
    self.assertEqual(conv._fill_shape(x, 3), (x, x, x))
    self.assertEqual(conv._fill_shape(x, 4), (x, x, x, x))
    self.assertEqual(conv._fill_shape([x, x + 1, x + 2], 3),
                     (x, x + 1, x + 2))

    err = "n must be a positive integer"
    with self.assertRaisesRegexp(TypeError, err):
      conv._fill_shape(x, 0)

    err = ("must be either a positive integer or an iterable of positive "
           "integers of size 4")
    with self.assertRaisesRegexp(TypeError, err):
      conv._fill_shape([], 4)
    with self.assertRaisesRegexp(TypeError, err):
      conv._fill_shape([x], 4)
    with self.assertRaisesRegexp(TypeError, err):
      conv._fill_shape([x, x], 4)
    with self.assertRaisesRegexp(TypeError, err):
      conv._fill_shape(["b"], 4)


class DefaultTransposeSizeTest(parameterized.TestCase, tf.test.TestCase):

  # Constants for use in parameterized test.
  input_shape = [[20], [23, 11, 13], [1, 3]]
  stride = [[3], [7, 1, 2], [6, 2]]
  kernel_shape = [[4], [1, 3, 2], [34, 2]]
  padding = [snt.SAME, snt.VALID, snt.VALID]
  output_shape = []

  for i, pad in enumerate(padding):
    if pad == snt.SAME:
      output_shape.append([x * y for x, y in zip(input_shape[i], stride[i])])
    if pad == snt.VALID:
      output_shape.append([x * y + z - 1 for x, y, z in
                           zip(input_shape[i], stride[i], kernel_shape[i])])

  @parameterized.parameters(
      *zip(input_shape, stride, kernel_shape, padding, output_shape))
  def testFunction(self, input_shape, stride, kernel_shape, padding,
                   output_shape):
    """Test output shapes are correct."""
    self.assertEqual(conv._default_transpose_size(input_shape, stride,
                                                  kernel_shape=kernel_shape,
                                                  padding=padding),
                     tuple(output_shape))

  @parameterized.parameters(
      *zip(input_shape, stride, kernel_shape, padding, output_shape))
  def testModules(self, input_shape, stride, kernel_shape, padding,
                  output_shape):
    """Test ConvTranspose modules return expected default output shapes."""
    if len(input_shape) == 1:
      module = snt.Conv1DTranspose
    elif len(input_shape) == 2:
      module = snt.Conv2DTranspose
    elif len(input_shape) == 3:
      module = snt.Conv3DTranspose

    batch_size = [1]
    channels = [1]

    inputs = tf.zeros(shape=batch_size + input_shape + channels,
                      dtype=tf.float32)
    outputs = module(output_channels=1, kernel_shape=kernel_shape,
                     stride=stride, padding=padding)(inputs)
    self.assertEqual(output_shape, outputs.get_shape().as_list()[1:-1])

  @parameterized.parameters(
      *zip(input_shape, stride, kernel_shape, padding, output_shape))
  def testConnectTwice(self, input_shape, stride, kernel_shape, padding,
                       output_shape):
    """Test ConvTranspose modules with multiple connections."""
    if len(input_shape) == 1:
      module = snt.Conv1DTranspose
    elif len(input_shape) == 2:
      module = snt.Conv2DTranspose
    elif len(input_shape) == 3:
      module = snt.Conv3DTranspose

    batch_size = [1]
    channels = [1]

    inputs = tf.zeros(shape=batch_size + input_shape + channels,
                      dtype=tf.float32)
    inputs_2 = tf.zeros(shape=batch_size + input_shape + channels,
                        dtype=tf.float32)
    conv1 = module(output_channels=1, kernel_shape=kernel_shape,
                   stride=stride, padding=padding)
    outputs = conv1(inputs)

    # Connecting for the second time with the same shape should be OK.
    outputs_2 = conv1(inputs_2)

    # So should connecting with a different shape.
    new_input_shape = [25] * len(input_shape)
    new_inputs = tf.zeros(shape=batch_size + new_input_shape + channels,
                          dtype=tf.float32)
    new_outputs = conv1(new_inputs)

    with self.test_session() as sess:
      tf.global_variables_initializer().run()
      outputs_array, outputs_array_2 = sess.run([outputs, outputs_2])
      self.assertEqual(outputs_array.shape, outputs_array_2.shape)

      sess.run(new_outputs)


class SharedConvTest(parameterized.TestCase, tf.test.TestCase):

  CONV_1D_KWARGS = {
      "output_channels": 1,
      "kernel_shape": 3,
  }
  CONV_1D_MASKED_KWARGS = {
      "output_channels": 1,
      "kernel_shape": 3,
      "mask": np.zeros((3, 10, 1), dtype=np.float64),
  }
  CONV_1D_CAUSAL_KWARGS = {
      "output_channels": 1,
      "kernel_shape": 3,
      "padding": snt.CAUSAL,
  }
  SEPARABLE_CONV_1D_KWARGS = {
      "output_channels": 10,
      "channel_multiplier": 1,
      "kernel_shape": 3,
      "rate": (2,),
  }
  CONV_2D_KWARGS = CONV_1D_KWARGS
  CONV_3D_KWARGS = CONV_1D_KWARGS
  CONV_3D_MIXED_PADDING_KWARGS = {
      "output_channels": 1,
      "kernel_shape": 3,
      "padding": [snt.CAUSAL, snt.SAME, snt.FULL]
  }
  DEPTHWISE_CONV_2D_KWARGS = {
      "channel_multiplier": 1,
      "kernel_shape": 3,
  }
  SEPARABLE_CONV_2D_KWARGS = {
      "output_channels": 10,
      "channel_multiplier": 1,
      "kernel_shape": 3,
      "rate": (2, 1),
  }
  IN_PLANE_CONV_2D_KWARGS = {
      "kernel_shape": 3,
  }
  CONV_1D_TRANSPOSE_KWARGS = {
      "output_channels": 1,
      "output_shape": [10],
      "kernel_shape": 3,
  }
  CONV_2D_TRANSPOSE_KWARGS = {
      "output_channels": 1,
      "output_shape": [10, 10],
      "kernel_shape": 3,
  }
  CONV_3D_TRANSPOSE_KWARGS = {
      "output_channels": 1,
      "output_shape": [10, 10, 10],
      "kernel_shape": 3,
  }

  modules = [
      (snt.Conv1D, 1, CONV_1D_KWARGS),
      (snt.Conv1D, 1, CONV_1D_MASKED_KWARGS),
      (snt.Conv1D, 1, CONV_1D_CAUSAL_KWARGS),
      (snt.Conv2D, 2, CONV_2D_KWARGS),
      (snt.Conv3D, 3, CONV_3D_KWARGS),
      (snt.Conv3D, 3, CONV_3D_MIXED_PADDING_KWARGS),
      (snt.Conv1DTranspose, 1, CONV_1D_TRANSPOSE_KWARGS),
      (snt.Conv2DTranspose, 2, CONV_2D_TRANSPOSE_KWARGS),
      (snt.Conv3DTranspose, 3, CONV_3D_TRANSPOSE_KWARGS),
      (snt.DepthwiseConv2D, 2, DEPTHWISE_CONV_2D_KWARGS),
      (snt.InPlaneConv2D, 2, IN_PLANE_CONV_2D_KWARGS),
      (snt.SeparableConv1D, 1, SEPARABLE_CONV_1D_KWARGS),
      (snt.SeparableConv2D, 2, SEPARABLE_CONV_2D_KWARGS),
  ]

  @parameterized.parameters(*modules)
  def testPartitioners(self, module, num_input_dims, module_kwargs):
    inputs = tf.zeros((10,) * (num_input_dims + 2))

    keys = module.get_possible_initializer_keys(use_bias=True)
    partitioners = {
        key: tf.variable_axis_size_partitioner(10) for key in keys
    }

    if "mask" in module_kwargs:
      np_dtype = inputs.dtype.as_numpy_dtype
      module_kwargs["mask"] = module_kwargs["mask"].astype(np_dtype)

    convolution = module(partitioners=partitioners, **module_kwargs)
    convolution(inputs)

    for key in keys:
      self.assertEqual(type(getattr(convolution, key)),
                       variables.PartitionedVariable)

    try:
      convolution_t = convolution.transpose()
    except (AttributeError, snt.NotSupportedError):
      return

    self.assertEqual(convolution_t.partitioners, convolution.partitioners)

  @parameterized.parameters(
      *itertools.product(modules, (True, False),
                         (tf.float16, tf.float32, tf.float64)))
  def testVariables(self, module_info, use_bias, dtype):
    """The correct number of variables are created."""
    module, num_input_dims, module_kwargs = module_info

    mod_name = "module"

    input_shape = (10,) * (num_input_dims + 2)
    inputs = tf.placeholder(dtype, input_shape)

    if "mask" in module_kwargs:
      np_dtype = dtype.as_numpy_dtype
      module_kwargs["mask"] = module_kwargs["mask"].astype(np_dtype)

    with tf.variable_scope("scope"):
      conv_mod = module(name=mod_name, use_bias=use_bias, **module_kwargs)

    self.assertEqual(conv_mod.scope_name, "scope/" + mod_name)
    self.assertEqual(conv_mod.module_name, mod_name)

    with self.assertRaisesRegexp(snt.NotConnectedError, "not instantiated yet"):
      conv_mod.get_variables()

    output = conv_mod(inputs)

    self.assertEqual(dtype, output.dtype)

    # Check that the graph and module has the correct number of variables: one
    # two, or three, depending on module and configuration.
    supposed_variables = conv_mod.get_possible_initializer_keys(
        use_bias=use_bias)
    self.assertIn(len(supposed_variables), [1, 2, 3])

    graph_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
    self.assertEqual(len(graph_variables), len(supposed_variables))
    conv_variables = conv_mod.get_variables()
    self.assertEqual(len(conv_variables), len(supposed_variables))

    variable_names = {v.name for v in conv_variables}

    for var_name in supposed_variables:
      self.assertIn("scope/{}/{}:0".format(mod_name, var_name), variable_names)

    with self.test_session() as sess:
      tf.global_variables_initializer().run()
      inputs_data = np.random.rand(*input_shape)
      sess.run(output, feed_dict={inputs: inputs_data})

  @parameterized.parameters(*itertools.product(modules, (True, False)))
  def testMissingChannelsError(self, module_info, use_bias):
    """Error is thrown if the input is missing a channel dimension."""
    module, num_input_dims, module_kwargs = module_info
    conv_mod = module(use_bias=use_bias, **module_kwargs)

    inputs = tf.placeholder(tf.float32, (10,) * (num_input_dims + 1))

    err = "Input Tensor must have"
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      conv_mod(inputs)

  @parameterized.parameters(*itertools.product(modules, (True, False)))
  def testNonDefinedChannelsDimension(self, module_info, use_bias):
    """Error is thrown if the input's channel dimension isn't defined."""
    module, num_input_dims, module_kwargs = module_info
    conv_mod = module(use_bias=use_bias, **module_kwargs)

    inputs = tf.placeholder(tf.float32, (10,) * (num_input_dims + 1) + (None,))
    err = "Number of input channels"
    with self.assertRaisesRegexp(snt.UnderspecifiedError, err):
      conv_mod(inputs)

  @parameterized.parameters(*itertools.product(modules, (True, False)))
  def testFlattenedError(self, module_info, use_bias):
    """Error is thrown if the input has been incorrectly flattened."""
    module, num_input_dims, module_kwargs = module_info
    conv_mod = module(use_bias=use_bias, **module_kwargs)

    inputs = tf.placeholder(tf.float32, (10,) * (num_input_dims + 1))
    inputs = snt.BatchFlatten()(inputs)

    err = "Input Tensor must have"
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      conv_mod(inputs)

  @parameterized.parameters(*modules)
  def testCustomGetter(self, module, num_input_dims, module_kwargs):
    """Check that custom_getter option works."""

    def stop_gradient(getter, *args, **kwargs):
      return tf.stop_gradient(getter(*args, **kwargs))

    inputs = tf.placeholder(tf.float32, (10,) * (num_input_dims + 2))

    if "mask" in module_kwargs:
      np_dtype = inputs.dtype.as_numpy_dtype
      module_kwargs["mask"] = module_kwargs["mask"].astype(np_dtype)

    conv_mod1 = module(**module_kwargs)
    out1 = conv_mod1(inputs)

    conv_mod2 = module(custom_getter=stop_gradient, **module_kwargs)
    out2 = conv_mod2(inputs)

    num_variables = len(conv_mod1.get_variables())

    grads1 = tf.gradients(out1, list(conv_mod1.get_variables()))
    grads2 = tf.gradients(out2, list(conv_mod2.get_variables()))

    self.assertEqual([tf.Tensor] * num_variables, [type(g) for g in grads1])
    self.assertEqual([None] * num_variables, grads2)

    # Check that the transpose, if present, also adopts the custom getter.
    try:
      conv_mod2_transpose = conv_mod2.transpose()
    except (AttributeError, snt.NotSupportedError):
      return
    inputs_transpose = tf.placeholder(tf.float32, out2.get_shape())
    out3 = conv_mod2_transpose(inputs_transpose)
    grads3 = tf.gradients(out3, list(conv_mod2_transpose.get_variables()))
    self.assertEqual([None] * num_variables, grads3)


# These functions compute the expected output shape of a convolution of each
# padding type, for a given input shape and kernel size.
_PADDINGS_EXPECTED_SHAPE_TRANSFORMS = {
    snt.SAME: lambda in_length, kernel_size: in_length,
    snt.VALID: lambda in_length, kernel_size: in_length - kernel_size + 1,
    snt.FULL: lambda in_length, kernel_size: in_length + kernel_size - 1,
    snt.CAUSAL: lambda in_length, kernel_size: in_length,
    snt.REVERSE_CAUSAL: lambda in_length, kernel_size: in_length,
}


class Conv2DTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.parameters(*itertools.product(
      [True, False],  # use_bias
      conv.SUPPORTED_2D_DATA_FORMATS,  # data_format
      conv.ALLOWED_PADDINGS,  # padding_height
      conv.ALLOWED_PADDINGS,  # padding_width
      conv.ALLOWED_PADDING_VALUES,  # padding_value
  ))
  def testShapes(self, use_bias, data_format, padding_height, padding_width,
                 padding_value):
    """The generated shapes are correct with different paddings."""

    batch_size = random.randint(1, 100)
    in_height = random.randint(10, 288)
    in_width = random.randint(10, 288)
    in_channels = random.randint(1, 10)
    out_channels = random.randint(1, 32)
    kernel_shape_h = random.randint(1, 11)
    kernel_shape_w = random.randint(1, 11)

    if data_format == conv.DATA_FORMAT_NHWC:
      inputs = tf.placeholder(
          tf.float32, shape=[batch_size, in_height, in_width, in_channels])
    else:  # NCHW
      inputs = tf.placeholder(
          tf.float32, shape=[batch_size, in_channels, in_height, in_width])

    conv1 = snt.Conv2D(
        name="conv1",
        output_channels=out_channels,
        kernel_shape=[kernel_shape_h, kernel_shape_w],
        padding=[padding_height, padding_width],
        padding_value=padding_value,
        data_format=data_format,
        stride=1,
        use_bias=use_bias)

    output = conv1(inputs)
    expected_out_height = _PADDINGS_EXPECTED_SHAPE_TRANSFORMS[padding_height](
        in_height, kernel_shape_h)
    expected_out_width = _PADDINGS_EXPECTED_SHAPE_TRANSFORMS[padding_width](
        in_width, kernel_shape_w)

    if data_format == conv.DATA_FORMAT_NHWC:
      self.assertTrue(output.get_shape().is_compatible_with(
          [batch_size, expected_out_height, expected_out_width, out_channels]))
    else:  # NCHW
      self.assertTrue(output.get_shape().is_compatible_with(
          [batch_size, out_channels, expected_out_height, expected_out_width]))

    self.assertTrue(
        conv1.w.get_shape().is_compatible_with(
            [kernel_shape_h, kernel_shape_w, in_channels, out_channels]))

    if use_bias:
      self.assertTrue(
          conv1.b.get_shape().is_compatible_with(
              [out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesNotKnown(self, use_bias):
    """The generated shapes are correct when input shape not known."""

    batch_size = 5
    in_height = in_width = 32
    in_channels = out_channels = 5
    kernel_shape_h = kernel_shape_w = 3

    inputs = tf.placeholder(
        tf.float32,
        shape=[None, None, None, in_channels],
        name="inputs")

    conv1 = snt.Conv2D(
        name="conv1",
        output_channels=out_channels,
        kernel_shape=[kernel_shape_h, kernel_shape_w],
        padding=snt.SAME,
        stride=1,
        use_bias=use_bias)

    output = conv1(inputs)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      output_eval = output.eval({
          inputs: np.zeros([batch_size, in_height, in_width, in_channels])})

      self.assertEqual(
          output_eval.shape,
          (batch_size, in_height, in_width, out_channels))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesNotKnownAtrous(self, use_bias):
    """No error is thrown if image shape isn't known for atrous convolution."""

    inputs = tf.placeholder(
        tf.float32,
        shape=[None, None, None, 5],
        name="inputs")

    conv1 = snt.Conv2D(
        name="conv1",
        output_channels=5,
        kernel_shape=[3, 3],
        padding=snt.SAME,
        stride=1,
        rate=2,
        use_bias=use_bias)

    conv1(inputs)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testKernelShape(self, use_bias):
    """Errors are thrown for invalid kernel shapes."""

    snt.Conv2D(output_channels=10, kernel_shape=[3, 4], name="conv1",
               use_bias=use_bias)
    snt.Conv2D(output_channels=10, kernel_shape=3, name="conv1",
               use_bias=use_bias)

    err = "Invalid kernel shape"
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv2D(output_channels=10,
                 kernel_shape=[3, 3, 3],
                 name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testStrideError(self, use_bias):
    """Errors are thrown for invalid strides."""

    snt.Conv2D(
        output_channels=10, kernel_shape=3, stride=1, name="conv1",
        use_bias=use_bias)
    snt.Conv2D(
        output_channels=10, kernel_shape=3, stride=[1, 1], name="conv1",
        use_bias=use_bias)
    snt.Conv2D(
        output_channels=10, kernel_shape=3, stride=[1, 1, 1, 1], name="conv1",
        use_bias=use_bias)

    with self.assertRaisesRegexp(snt.IncompatibleShapeError, "Invalid stride"):
      snt.Conv2D(output_channels=10,
                 kernel_shape=3,
                 stride=[1, 1, 1],
                 name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRateError(self, use_bias):
    """Errors are thrown for invalid dilation rates."""

    snt.Conv2D(
        output_channels=10, kernel_shape=3, rate=1, name="conv1",
        use_bias=use_bias)
    snt.Conv2D(
        output_channels=10, kernel_shape=3, rate=2, name="conv1",
        use_bias=use_bias)

    for rate in [0, 0.5, -1]:
      with self.assertRaisesRegexp(snt.IncompatibleShapeError,
                                   "Invalid rate shape*"):
        snt.Conv2D(output_channels=10,
                   kernel_shape=3,
                   rate=rate,
                   name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRateAndStrideError(self, use_bias):
    """Errors are thrown for stride > 1 when using atrous convolution."""
    err = "Cannot have stride > 1 with rate > 1"
    with self.assertRaisesRegexp(snt.NotSupportedError, err):
      snt.Conv2D(output_channels=10, kernel_shape=3,
                 stride=2, rate=2, name="conv1", use_bias=use_bias)
    with self.assertRaisesRegexp(snt.NotSupportedError, err):
      snt.Conv2D(output_channels=10, kernel_shape=3,
                 stride=[2, 1], rate=2, name="conv1", use_bias=use_bias)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInputTypeError(self, use_bias):
    """Errors are thrown for invalid input types."""
    conv1 = snt.Conv2D(output_channels=1,
                       kernel_shape=3,
                       stride=1,
                       padding=snt.SAME,
                       name="conv1",
                       use_bias=use_bias,
                       initializers=create_constant_initializers(
                           1.0, 1.0, use_bias))

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([1, 5, 5, 1]), dtype=dtype)
      err = "Input must have dtype tf.float.*"
      with self.assertRaisesRegexp(TypeError, err):
        conv1(x)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInitializers(self, use_bias):
    """Test initializers work as expected."""
    w = random.random()
    b = random.random()

    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(w, b, use_bias))

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          conv1.w.eval(),
          np.full([3, 3, 2, 1], w, dtype=np.float32))

      if use_bias:
        self.assertAllClose(
            conv1.b.eval(),
            [b])

    err = "Initializer for 'w' is not a callable function or dictionary"
    with self.assertRaisesRegexp(TypeError, err):
      snt.Conv2D(output_channels=10, kernel_shape=3, stride=1, name="conv1",
                 initializers={"w": tf.ones([])})

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv1",
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRegularizersInRegularizationLosses(self, use_bias):
    regularizers = create_regularizers(use_bias,
                                       contrib_layers.l1_regularizer(scale=0.5))

    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        regularizers=regularizers,
        use_bias=use_bias,
        name="conv1")
    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    graph_regularizers = tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES)
    self.assertRegexpMatches(graph_regularizers[0].name, ".*l1_regularizer.*")
    if use_bias:
      self.assertRegexpMatches(graph_regularizers[1].name, ".*l1_regularizer.*")

  @parameterized.parameters(*itertools.product(
      [True, False],  # use_bias
      [snt.CONSTANT_PADDING],  # padding_value
      # (padding_h, padding_w, expected_out):
      [(snt.VALID, snt.VALID, [[9, 9, 9],
                               [9, 9, 9],
                               [9, 9, 9]]),
       (snt.SAME, snt.SAME, [[4, 6, 6, 6, 4],
                             [6, 9, 9, 9, 6],
                             [6, 9, 9, 9, 6],
                             [6, 9, 9, 9, 6],
                             [4, 6, 6, 6, 4]]),
       (snt.CAUSAL, snt.CAUSAL, [[1, 2, 3, 3, 3],
                                 [2, 4, 6, 6, 6],
                                 [3, 6, 9, 9, 9],
                                 [3, 6, 9, 9, 9],
                                 [3, 6, 9, 9, 9]]),
       (snt.REVERSE_CAUSAL, snt.REVERSE_CAUSAL, [[9, 9, 9, 6, 3],
                                                 [9, 9, 9, 6, 3],
                                                 [9, 9, 9, 6, 3],
                                                 [6, 6, 6, 4, 2],
                                                 [3, 3, 3, 2, 1]]),
       (snt.FULL, snt.FULL, [[1, 2, 3, 3, 3, 2, 1],
                             [2, 4, 6, 6, 6, 4, 2],
                             [3, 6, 9, 9, 9, 6, 3],
                             [3, 6, 9, 9, 9, 6, 3],
                             [3, 6, 9, 9, 9, 6, 3],
                             [2, 4, 6, 6, 6, 4, 2],
                             [1, 2, 3, 3, 3, 2, 1]]),
       (snt.CAUSAL, snt.FULL, [[1, 2, 3, 3, 3, 2, 1],
                               [2, 4, 6, 6, 6, 4, 2],
                               [3, 6, 9, 9, 9, 6, 3],
                               [3, 6, 9, 9, 9, 6, 3],
                               [3, 6, 9, 9, 9, 6, 3]]),
       (snt.SAME, snt.REVERSE_CAUSAL, [[6, 6, 6, 4, 2],
                                       [9, 9, 9, 6, 3],
                                       [9, 9, 9, 6, 3],
                                       [9, 9, 9, 6, 3],
                                       [6, 6, 6, 4, 2]])],
  ))
  def testComputation(self, use_bias, padding_value, padding_and_expected_out):
    """Run through for something with a known answer using different args."""
    padding_h, padding_w, expected_out = padding_and_expected_out
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=(padding_h, padding_w),
        padding_value=padding_value,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    out = tf.squeeze(out, axis=(0, 3))
    expected_out = np.asarray(expected_out, dtype=np.float32)
    if use_bias:
      expected_out += 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(out.eval(), expected_out)

  @parameterized.parameters(*itertools.product(
      # (padding_h, padding_w, padding_value, expected_out):
      # pylint:disable=bad-whitespace
      [(snt.VALID, snt.VALID, snt.CONSTANT_PADDING, [[26, 28, 23],
                                                     [29, 31, 26],
                                                     [25, 27, 29]]),
       (snt.VALID, snt.VALID, snt.REFLECT_PADDING, [[26, 28, 23],
                                                    [29, 31, 26],
                                                    [25, 27, 29]]),
       (snt.VALID, snt.VALID, snt.SYMMETRIC_PADDING, [[26, 28, 23],
                                                      [29, 31, 26],
                                                      [25, 27, 29]]),
       (snt.SAME, snt.SAME, snt.CONSTANT_PADDING, [[12, 14, 13, 12, 10],
                                                   [19, 26, 28, 23, 16],
                                                   [21, 29, 31, 26, 18],
                                                   [16, 25, 27, 29, 20],
                                                   [ 9, 13, 12, 18, 14]]),
       (snt.SAME, snt.SAME, snt.REFLECT_PADDING, [[36, 25, 20, 15, 18],
                                                  [30, 26, 28, 23, 26],
                                                  [33, 29, 31, 26, 29],
                                                  [22, 25, 27, 29, 32],
                                                  [16, 19, 21, 30, 33]]),
       (snt.SAME, snt.SAME, snt.SYMMETRIC_PADDING, [[18, 17, 19, 21, 27],
                                                    [27, 26, 28, 23, 22],
                                                    [30, 29, 31, 26, 25],
                                                    [26, 25, 27, 29, 28],
                                                    [28, 20, 15, 24, 30]]),
       (snt.REVERSE_CAUSAL, snt.REVERSE_CAUSAL, snt.CONSTANT_PADDING,
        [[26, 28, 23, 16, 6],
         [29, 31, 26, 18, 7],
         [25, 27, 29, 20, 8],
         [13, 12, 18, 14, 8],
         [ 7,  3,  6,  5, 3]]),
       (snt.REVERSE_CAUSAL, snt.REVERSE_CAUSAL, snt.REFLECT_PADDING,
        [[26, 28, 23, 26, 23],
         [29, 31, 26, 29, 26],
         [25, 27, 29, 32, 29],
         [19, 21, 30, 33, 30],
         [25, 27, 29, 32, 29]]),
       (snt.REVERSE_CAUSAL, snt.REVERSE_CAUSAL, snt.SYMMETRIC_PADDING,
        [[26, 28, 23, 22, 22],
         [29, 31, 26, 25, 25],
         [25, 27, 29, 28, 28],
         [20, 15, 24, 30, 30],
         [20, 15, 24, 30, 30]]),
       (snt.FULL, snt.FULL, snt.CONSTANT_PADDING,
        [[ 0,  1,  3,  6,  9,  7,  4],
         [ 5, 12, 14, 13, 12, 10,  6],
         [ 8, 19, 26, 28, 23, 16,  6],
         [ 9, 21, 29, 31, 26, 18,  7],
         [10, 16, 25, 27, 29, 20,  8],
         [ 7,  9, 13, 12, 18, 14,  8],
         [ 6,  6,  7,  3,  6,  5,  3]]),
       (snt.FULL, snt.FULL, snt.REFLECT_PADDING,
        [[26, 30, 26, 28, 23, 26, 23],
         [25, 36, 25, 20, 15, 18, 15],
         [26, 30, 26, 28, 23, 26, 23],
         [29, 33, 29, 31, 26, 29, 26],
         [25, 22, 25, 27, 29, 32, 29],
         [19, 16, 19, 21, 30, 33, 30],
         [25, 22, 25, 27, 29, 32, 29]]),
       (snt.FULL, snt.FULL, snt.SYMMETRIC_PADDING,
        [[18, 18, 17, 19, 21, 27, 27],
         [18, 18, 17, 19, 21, 27, 27],
         [27, 27, 26, 28, 23, 22, 22],
         [30, 30, 29, 31, 26, 25, 25],
         [26, 26, 25, 27, 29, 28, 28],
         [28, 28, 20, 15, 24, 30, 30],
         [28, 28, 20, 15, 24, 30, 30]]),
       (snt.CAUSAL, snt.CAUSAL, snt.CONSTANT_PADDING, [[ 0,  1,  3,  6,  9],
                                                       [ 5, 12, 14, 13, 12],
                                                       [ 8, 19, 26, 28, 23],
                                                       [ 9, 21, 29, 31, 26],
                                                       [10, 16, 25, 27, 29]]),
       (snt.CAUSAL, snt.CAUSAL, snt.REFLECT_PADDING, [[26, 30, 26, 28, 23],
                                                      [25, 36, 25, 20, 15],
                                                      [26, 30, 26, 28, 23],
                                                      [29, 33, 29, 31, 26],
                                                      [25, 22, 25, 27, 29]]),
       (snt.CAUSAL, snt.CAUSAL, snt.SYMMETRIC_PADDING, [[18, 18, 17, 19, 21],
                                                        [18, 18, 17, 19, 21],
                                                        [27, 27, 26, 28, 23],
                                                        [30, 30, 29, 31, 26],
                                                        [26, 26, 25, 27, 29]]),
      ],
  ))
  def testPaddingValues(self, padding_and_expected_out):
    """Run through for something with a known answer using different args."""
    padding_h, padding_w, padding_value, expected_out = padding_and_expected_out
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=(padding_h, padding_w),
        padding_value=padding_value,
        name="conv1",
        use_bias=False,
        initializers=create_constant_initializers(1.0, 1.0, False))

    # Make a less trivial input test tensor.
    in_tensor = np.arange(5 * 5).reshape([1, 5, 5, 1])
    # Do modulo 7 to not have large output values.
    in_tensor = in_tensor % 7

    out = conv1(tf.constant(in_tensor, dtype=np.float32))
    out = tf.squeeze(out, axis=(0, 3))

    expected_out = np.asarray(expected_out, dtype=np.float32)

    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""

    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=snt.SAME,
        use_bias=use_bias,
        name="conv1")

    x = np.random.randn(1, 5, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

      # Now change the weights
      w = np.random.randn(3, 3, 1, 1)
      conv1.w.assign(w).eval()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testAtrousConvValid(self, use_bias):
    """The atrous conv is constructed and applied correctly with snt.VALID."""
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        rate=2,
        padding=snt.VALID,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 0.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(np.reshape(out.eval(), [1, 1]), [[9]])

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testAtrousConvSame(self, use_bias):
    """The atrous conv 2D is constructed and applied correctly with SAME."""
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        rate=2,
        padding=snt.SAME,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    expected_out = np.array([[5, 5, 7, 5, 5],
                             [5, 5, 7, 5, 5],
                             [7, 7, 10, 7, 7],
                             [5, 5, 7, 5, 5],
                             [5, 5, 7, 5, 5]])
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(np.reshape(out.eval(), [5, 5]), expected_out)

  def testClone(self):
    net = snt.Conv2D(name="conv2d",
                     output_channels=4,
                     kernel_shape=3,
                     stride=5)
    clone1 = net.clone()
    clone2 = net.clone(name="clone2")

    input_to_net = tf.placeholder(tf.float32, shape=[None, 100, 100, 3])
    net_out = net(input_to_net)
    clone1_out = clone1(input_to_net)
    clone2_out = clone2(input_to_net)

    all_vars = tf.trainable_variables()
    net_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=net.variable_scope.name + "/")
    clone1_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone1.variable_scope.name + "/")
    clone2_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone2.variable_scope.name + "/")

    self.assertEqual(net.output_channels, clone1.output_channels)
    self.assertEqual(net.module_name + "_clone", clone1.module_name)
    self.assertEqual("clone2", clone2.module_name)
    self.assertLen(all_vars, 3*len(net_vars))
    self.assertLen(net_vars, len(clone1_vars))
    self.assertLen(net_vars, len(clone2_vars))
    self.assertEqual(net_out.get_shape().as_list(),
                     clone1_out.get_shape().as_list())
    self.assertEqual(net_out.get_shape().as_list(),
                     clone2_out.get_shape().as_list())

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testTransposition(self, use_bias):
    """Tests if the correct output shapes are setup in transposed module."""
    net = snt.Conv2D(name="conv2d",
                     output_channels=4,
                     kernel_shape=3,
                     stride=1,
                     use_bias=use_bias)

    net_transpose = net.transpose()
    input_to_net = tf.placeholder(tf.float32, shape=[None, 100, 100, 3])

    err = "Variables in {} not instantiated yet, __call__ the module first."
    with self.assertRaisesRegexp(snt.NotConnectedError,
                                 err.format(net.scope_name)):
      net_transpose(input_to_net)
    net_transpose = net.transpose(name="another_net_transpose")
    net_out = net(input_to_net)
    net_transposed_output = net_transpose(net_out)
    self.assertAllEqual(net_transposed_output.get_shape().as_list(),
                        input_to_net.get_shape().as_list())

  def testMask1D(self):
    """1D Masks are applied properly."""
    mask = np.ones((3,), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(1, 5, 5, 2))
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 18 * np.ones((1, 3, 3, 1))
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask2D(self):
    """2D Masks are applied properly."""

    # This mask, applied on an image filled with 1, should result in an image
    # filled with 8 (since we sum 4 elements per channel and there are 2 input
    # channels).
    mask = np.array([[1, 1, 1],
                     [1, 0, 0],
                     [0, 0, 0]], dtype=np.float32)
    inputs = tf.constant(1.0, shape=(1, 5, 5, 2))
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = np.array([[8] * 3] * 3)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(np.reshape(out.eval(), [3, 3]), expected_out)

  def testMask3D(self):
    """3D Masks are applied properly."""
    mask = np.ones((3, 3, 2), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(1, 5, 5, 2))
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 18 * np.ones((1, 3, 3, 1))
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask4D(self):
    """4D Masks are applied properly."""

    # This mask, applied on an image filled with 1, should result in an image
    # filled with 17, as there are 18 weights but we zero out one of them.
    mask = np.ones([3, 3, 2, 1], dtype=np.float32)
    mask[0, 0, 0, :] = 0
    inputs = tf.constant(1.0, shape=(1, 5, 5, 2))
    conv1 = snt.Conv2D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = np.array([[17] * 3] * 3)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(np.reshape(out.eval(), [3, 3]), expected_out)

  def testMaskErrorIncompatibleRank1(self):
    """Errors are thrown for incompatible rank 1 mask."""

    np_mask = np.ones((3,))
    x = tf.constant(0.0, shape=(2, 8, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv2D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank2(self):
    """Errors are thrown for incompatible rank 2 mask."""

    np_mask = np.ones((3, 3))
    x = tf.constant(0.0, shape=(2, 8, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv2D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank3(self):
    """Errors are thrown for incompatible rank 3 mask."""

    np_mask = np.ones((2, 4, 4))
    x = tf.constant(0.0, shape=(2, 8, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv2D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank4(self):
    """Errors are thrown for incompatible rank 4 mask."""

    np_mask = np.ones((3, 3, 4, 5))
    x = tf.constant(0.0, shape=(2, 8, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv2D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIfIncorrectDtype(self):
    """Errors are thrown when a Tensor with incorrect dtype is used."""

    mask = tf.constant(0, shape=(4, 4), dtype=tf.int32)
    x = tf.constant(0.0, shape=(2, 8, 8, 6))

    with self.assertRaises(TypeError) as cm:
      snt.Conv2D(output_channels=4, kernel_shape=(4, 4), mask=mask)(x)
    self.assertTrue(str(cm.exception).startswith(
        "Mask needs to have dtype float16, bfloat16, float32 or float64"))

  def testDataFormatNotSupported(self):
    """Errors are thrown when an unsupported data_format is used."""

    x = tf.constant(0.0, shape=(2, 8, 8, 6))
    data_format = "NWCH"
    self.assertNotIn(data_format, conv.SUPPORTED_2D_DATA_FORMATS)

    with self.assertRaisesRegexp(ValueError, "Invalid data_format"):
      snt.Conv2D(output_channels=4, kernel_shape=(4, 4),
                 data_format=data_format)(x)


class Conv2DTransposeTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    """Set up some variables to re-use in multiple tests."""

    super(Conv2DTransposeTest, self).setUp()

    self.batch_size = 100
    self.in_height = 32
    self.in_width = 32
    self.in_channels = 3
    self.out_channels = 10
    self.kernel_shape_h = 5
    self.kernel_shape_w = 5
    self.strides = (1, 1, 1, 1)
    self.padding = snt.SAME

    self.in_shape = (self.batch_size, self.in_height, self.in_width,
                     self.in_channels)

    self.out_shape = (self.in_height, self.in_width)

    self.kernel_shape = (self.kernel_shape_h, self.kernel_shape_w)

    self.kernel_shape2 = (self.kernel_shape_h, self.kernel_shape_w,
                          self.out_channels, self.in_channels)

  def testKernelsNotSpecified(self):
    """Tests error is raised if kernel shape is not specified."""
    with self.assertRaisesRegexp(ValueError, "`kernel_shape` cannot be None."):
      snt.Conv2DTranspose(output_channels=1)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testOutputShapeConsistency(self, use_bias):
    """Tests if output shapes are valid."""

    # When padding is SAME, then the actual number of padding pixels can be
    # computed as: pad = kernel_shape - strides + (-input_shape % strides)
    #                 =     5         -    1    + (- 32       %      1) = 4

    # The formula for the minimal size is:
    # oH = strides[1] * (in_height - 1) - padding + kernel_shape_h
    # oH =          1 * (       32 - 1) -    4    +       5 = 32

    # The formula for the maximum size (due to extra pixels) is:
    # oH_max = oH + strides[1] - 1
    # so, for strides = 1 and padding = SAME, input size == output size.
    inputs = tf.placeholder(tf.float32, shape=self.in_shape)

    conv1 = snt.Conv2DTranspose(name="conv2d_1",
                                output_channels=self.out_channels,
                                output_shape=self.out_shape,
                                kernel_shape=self.kernel_shape,
                                padding=self.padding,
                                stride=1,
                                use_bias=use_bias)

    outputs = conv1(inputs)

    self.assertTrue(outputs.get_shape().is_compatible_with((
        self.batch_size,) + self.out_shape + (self.out_channels,)))

    self.assertTrue(conv1.w.get_shape().is_compatible_with(self.kernel_shape2))
    if use_bias:
      self.assertTrue(conv1.b.get_shape().is_compatible_with(
          [self.out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testOutputShapeInteger(self, use_bias):
    """Tests if output shapes are valid when specified as an integer."""
    inputs = tf.zeros(shape=[3, 5, 5, 2], dtype=tf.float32)
    inputs_2 = tf.zeros(shape=[3, 5, 7, 2], dtype=tf.float32)

    conv1 = snt.Conv2DTranspose(name="conv2d_1",
                                output_channels=10,
                                output_shape=tf.Dimension(10),
                                kernel_shape=5,
                                padding=snt.SAME,
                                stride=2,
                                use_bias=use_bias)

    outputs = conv1(inputs)
    outputs_2 = conv1(inputs_2)

    self.assertTrue(outputs.get_shape().is_compatible_with((3, 10, 10, 10)))

    with self.test_session() as sess:
      tf.global_variables_initializer().run()
      sess.run(outputs)
      with self.assertRaises(tf.errors.InvalidArgumentError):
        sess.run(outputs_2)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testTransposition(self, use_bias):
    """Tests if the correct output shapes are setup in transposed module."""
    net = snt.Conv2DTranspose(name="conv2d",
                              output_channels=self.out_channels,
                              output_shape=self.out_shape,
                              kernel_shape=self.kernel_shape,
                              padding=self.padding,
                              stride=1,
                              use_bias=use_bias)

    net_transpose = net.transpose()
    input_to_net = tf.placeholder(tf.float32, shape=self.in_shape)
    err = "Variables in {} not instantiated yet, __call__ the module first."
    with self.assertRaisesRegexp(snt.NotConnectedError,
                                 err.format(net.scope_name)):
      net_transpose(input_to_net)
    net_transpose = net.transpose(name="another_net_transpose")
    net_out = net(input_to_net)
    net_transposed_output = net_transpose(net_out)
    self.assertEqual(net_transposed_output.get_shape(),
                     input_to_net.get_shape())

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.Conv2DTranspose(
        output_shape=(10, 10),
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv2d",
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  @parameterized.named_parameters(
      ("WithBiasWithOutputShape", True, True),
      ("WithBiasWithoutOutputShape", True, False),
      ("WithoutBiasWithOutputShape", False, True),
      ("WithoutBiasWithoutOutputShape", False, False))
  def testTransposeNHWC(self, use_bias, use_output_shape):
    """Test transpose for NHWC format."""
    output_shape = tf.TensorShape((4, 5))

    conv2_transpose = snt.Conv2DTranspose(
        output_channels=5,
        output_shape=output_shape if use_output_shape else None,
        kernel_shape=3,
        padding=snt.VALID,
        stride=1,
        name="conv2_transpose",
        use_bias=use_bias,
        data_format=conv.DATA_FORMAT_NHWC)
    conv2 = conv2_transpose.transpose()

    # Check kernel shapes, strides and padding match.
    self.assertEqual(conv2_transpose.kernel_shape, conv2.kernel_shape)
    self.assertEqual((1,) + conv2_transpose.stride[1:3] + (1,), conv2.stride)
    self.assertEqual(conv2_transpose.conv_op_padding, conv2.conv_op_padding)

    # Before conv2_transpose is connected, we cannot know how many
    # `output_channels` conv1 should have.
    err = "Variables in conv2_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv2.output_channels

    # After connection the number of `output_channels` is known.
    batch_size = 32
    in_height = 2
    in_width = 3
    in_channels = 4
    x = tf.constant(np.random.randn(batch_size, in_height, in_width,
                                    in_channels),
                    dtype=np.float32)
    conv2_transpose(x)
    self.assertEqual(in_channels, conv2.output_channels)

    # As is `output_channels`.
    self.assertEqual(output_shape, conv2_transpose.output_shape)

    # However, even after connection, the `input_shape` of the forward
    # convolution is not known until it is itself connected (i.e. it can be
    # connected to a different shape input from the `output_shape` of the
    # transpose convolution!)
    err = "Variables in conv2_transpose_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv2.input_shape

  @parameterized.named_parameters(
      ("WithBiasWithOutputShape", True, True),
      ("WithBiasWithoutOutputShape", True, False),
      ("WithoutBiasWithOutputShape", False, True),
      ("WithoutBiasWithoutOutputShape", False, False))
  def testTransposeNCHW(self, use_bias, use_output_shape):
    """Test transpose for NCHW format."""
    output_shape = tf.TensorShape((4, 5))

    conv2_transpose = snt.Conv2DTranspose(
        output_channels=5,
        output_shape=output_shape if use_output_shape else None,
        kernel_shape=3,
        padding=snt.VALID,
        stride=1,
        name="conv2_transpose",
        use_bias=use_bias,
        data_format=conv.DATA_FORMAT_NCHW)
    conv2 = conv2_transpose.transpose()

    # Check kernel shapes, strides and padding match.
    self.assertEqual(conv2_transpose.kernel_shape, conv2.kernel_shape)
    self.assertEqual((1,) + conv2_transpose.stride[1:3] + (1,), conv2.stride)
    self.assertEqual(conv2_transpose.conv_op_padding, conv2.conv_op_padding)

    # Before conv2_transpose is connected, we cannot know how many
    # `output_channels` conv1 should have.
    err = "Variables in conv2_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv2.output_channels

    # After connection the number of `output_channels` is known.
    batch_size = 32
    in_height = 2
    in_width = 3
    in_channels = 4
    x = tf.constant(np.random.randn(batch_size, in_channels, in_height,
                                    in_width),
                    dtype=np.float32)
    conv2_transpose(x)
    self.assertEqual(in_channels, conv2.output_channels)

    # As is `output_channels`.
    self.assertEqual(output_shape, conv2_transpose.output_shape)

    # However, even after connection, the `input_shape` of the forward
    # convolution is not known until it is itself connected (i.e. it can be
    # connected to a different shape input from the `output_shape` of the
    # transpose convolution!)
    err = "Variables in conv2_transpose_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv2.input_shape


class Conv1DTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.parameters(*itertools.product(
      [True, False],  # use_bias
      conv.SUPPORTED_1D_DATA_FORMATS,  # data_format
      conv.ALLOWED_PADDINGS  # padding
  ))
  def testShapes(self, use_bias, data_format, padding):
    """The generated shapes are correct with SAME and VALID padding."""

    batch_size = random.randint(1, 100)
    in_length = random.randint(10, 288)
    in_channels = random.randint(1, 10)
    out_channels = random.randint(1, 32)

    kernel_shape = random.randint(1, 10)

    if data_format == conv.DATA_FORMAT_NWC:
      inputs = tf.placeholder(
          tf.float32, shape=[batch_size, in_length, in_channels])
    else:  # NCW
      inputs = tf.placeholder(
          tf.float32, shape=[batch_size, in_channels, in_length])

    conv1 = snt.Conv1D(
        output_channels=out_channels,
        kernel_shape=kernel_shape,
        padding=padding,
        stride=1,
        data_format=data_format,
        name="conv1",
        use_bias=use_bias)

    output = conv1(inputs)

    expected_out_length = _PADDINGS_EXPECTED_SHAPE_TRANSFORMS[padding](
        in_length, kernel_shape)
    if data_format == conv.DATA_FORMAT_NWC:
      self.assertTrue(
          output.get_shape().is_compatible_with(
              [batch_size, expected_out_length, out_channels]))
    else:  # NCW
      self.assertTrue(
          output.get_shape().is_compatible_with(
              [batch_size, out_channels, expected_out_length]))

    self.assertTrue(
        conv1.w.get_shape().is_compatible_with(
            [kernel_shape, in_channels, out_channels]))

    if use_bias:
      self.assertTrue(
          conv1.b.get_shape().is_compatible_with(
              [out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesNotKnown(self, use_bias):
    """The generated shapes are correct when input shape not known."""

    batch_size = 5
    in_length = 32
    in_channels = out_channels = 5
    kernel_shape = 3

    inputs = tf.placeholder(
        tf.float32,
        shape=[None, None, in_channels],
        name="inputs")

    conv1 = snt.Conv1D(
        name="conv1",
        output_channels=out_channels,
        kernel_shape=kernel_shape,
        padding=snt.SAME,
        stride=1,
        use_bias=use_bias)

    output = conv1(inputs)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      output_eval = output.eval({
          inputs: np.zeros([batch_size, in_length, in_channels])})

      self.assertEqual(
          output_eval.shape,
          (batch_size, in_length, out_channels))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testKernelShape(self, use_bias):
    """Errors are thrown for invalid kernel shapes."""

    snt.Conv1D(output_channels=10, kernel_shape=[3], name="conv1",
               use_bias=use_bias)
    snt.Conv1D(output_channels=10, kernel_shape=3, name="conv1",
               use_bias=use_bias)

    err = "Invalid kernel shape"
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1D(output_channels=10, kernel_shape=[3, 3], name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testStrideError(self, use_bias):
    """Errors are thrown for invalid strides."""

    snt.Conv1D(
        output_channels=10, kernel_shape=3, stride=1, name="conv1",
        use_bias=use_bias)

    err = "Invalid stride"
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1D(output_channels=10, kernel_shape=3,
                 stride=[1, 1], name="conv1")

    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1D(output_channels=10, kernel_shape=3,
                 stride=[1, 1, 1, 1], name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRateError(self, use_bias):
    """Errors are thrown for invalid dilation rates."""

    snt.Conv1D(
        output_channels=10, kernel_shape=3, rate=1, name="conv1",
        use_bias=use_bias)
    snt.Conv1D(
        output_channels=10, kernel_shape=3, rate=2, name="conv1",
        use_bias=use_bias)

    for rate in [0, 0.5, -1]:
      with self.assertRaisesRegexp(snt.IncompatibleShapeError,
                                   "Invalid rate shape*"):
        snt.Conv1D(output_channels=10,
                   kernel_shape=3,
                   rate=rate,
                   name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRateAndStrideError(self, use_bias):
    """Errors are thrown for stride > 1 when using atrous convolution."""
    err = "Cannot have stride > 1 with rate > 1"
    with self.assertRaisesRegexp(snt.NotSupportedError, err):
      snt.Conv1D(output_channels=10, kernel_shape=3,
                 stride=2, rate=2, name="conv1", use_bias=use_bias)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInputTypeError(self, use_bias):
    """Errors are thrown for invalid input types."""
    conv1 = snt.Conv1D(output_channels=1,
                       kernel_shape=3,
                       stride=1,
                       padding=snt.VALID,
                       use_bias=use_bias,
                       name="conv1",
                       initializers=create_constant_initializers(
                           1.0, 1.0, use_bias))

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([1, 5, 1]), dtype=dtype)
      err = "Input must have dtype tf.float.*"
      with self.assertRaisesRegexp(TypeError, err):
        conv1(x)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInitializers(self, use_bias):
    """Test initializers work as expected."""
    w = random.random()
    b = random.random()

    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=snt.SAME,
        use_bias=use_bias,
        name="conv1",
        initializers=create_constant_initializers(w, b, use_bias))

    conv1(tf.placeholder(tf.float32, [1, 10, 2]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          conv1.w.eval(),
          np.full([3, 2, 1], w, dtype=np.float32))

      if use_bias:
        self.assertAllClose(
            conv1.b.eval(),
            [b])

    err = "Initializer for 'w' is not a callable function or dictionary"
    with self.assertRaisesRegexp(TypeError, err):
      snt.Conv1D(output_channels=10,
                 kernel_shape=3,
                 stride=1,
                 padding=snt.SAME,
                 use_bias=use_bias,
                 name="conv1",
                 initializers={"w": tf.ones([])})

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv1",
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [1, 10, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRegularizersInRegularizationLosses(self, use_bias):
    regularizers = create_regularizers(use_bias,
                                       contrib_layers.l1_regularizer(scale=0.5))

    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        regularizers=regularizers,
        name="conv1")
    conv1(tf.placeholder(tf.float32, [1, 10, 2]))

    graph_regularizers = tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES)
    self.assertRegexpMatches(graph_regularizers[0].name, ".*l1_regularizer.*")
    if use_bias:
      self.assertRegexpMatches(graph_regularizers[1].name, ".*l1_regularizer.*")

  @parameterized.parameters(*itertools.product(
      [True, False],  # use_bias
      [(snt.VALID, [3, 3, 3]),
       (snt.SAME, [2, 3, 3, 3, 2]),
       (snt.CAUSAL, [1, 2, 3, 3, 3]),
       (snt.FULL, [1, 2, 3, 3, 3, 2, 1]),
       (snt.REVERSE_CAUSAL, [3, 3, 3, 2, 1])]  # (padding, expected_out)
  ))
  def testComputation(self, use_bias, padding_and_expected_out):
    """Run through for something with a known answer using different args."""
    padding, expected_out = padding_and_expected_out

    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=padding,
        use_bias=use_bias,
        name="conv1",
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 1], dtype=np.float32)))
    expected_out = np.asarray(expected_out, dtype=np.float32)
    if use_bias:
      expected_out += 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(out.eval().flatten(), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""

    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=snt.SAME,
        use_bias=use_bias,
        name="conv1")

    x = np.random.randn(1, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

      # Now change the weights
      w = np.random.randn(3, 1, 1)
      conv1.w.assign(w).eval()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

  def testMask1D(self):
    """1D Masks are applied properly."""
    mask = np.ones((3,), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(1, 5, 2))
    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = np.reshape(np.array([6, 6, 6]), (1, 3, 1))
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask2D(self):
    """2D Masks are applied properly."""
    mask = np.ones((3, 2), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(1, 5, 2))
    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = np.reshape(np.array([6, 6, 6]), (1, 3, 1))
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask3D(self):
    """3D Masks are applied properly."""
    mask = np.ones((3, 2, 1), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(1, 5, 2))
    conv1 = snt.Conv1D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = np.reshape(np.array([6, 6, 6]), (1, 3, 1))
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMaskErrorIncompatibleRank1(self):
    """Errors are thrown for incompatible rank 1 mask."""

    np_mask = np.ones((3,))
    x = tf.constant(0.0, shape=(2, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv1D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank2(self):
    """Errors are thrown for incompatible rank 2 mask."""

    np_mask = np.ones((3, 3))
    x = tf.constant(0.0, shape=(2, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv1D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank3(self):
    """Errors are thrown for incompatible rank 3 mask."""

    np_mask = np.ones((2, 4, 4))
    x = tf.constant(0.0, shape=(2, 8, 6))

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv1D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIfIncorrectDtype(self):
    """Errors are thrown when a Tensor with incorrect dtype is used."""

    mask = tf.constant(0, shape=(4, 4), dtype=tf.int32)
    x = tf.constant(0.0, shape=(2, 8, 6))

    with self.assertRaises(TypeError) as cm:
      snt.Conv1D(output_channels=4, kernel_shape=4, mask=mask)(x)
    self.assertTrue(str(cm.exception).startswith(
        "Mask needs to have dtype float16, bfloat16, float32 or float64"))

  def testClone(self):
    net = snt.Conv1D(name="conv1d",
                     output_channels=4,
                     kernel_shape=3,
                     stride=5)
    clone1 = net.clone()
    clone2 = net.clone(name="clone2")

    input_to_net = tf.placeholder(tf.float32, shape=[None, 100, 3])
    net_out = net(input_to_net)
    clone1_out = clone1(input_to_net)
    clone2_out = clone2(input_to_net)

    all_vars = tf.trainable_variables()
    net_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=net.variable_scope.name + "/")
    clone1_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone1.variable_scope.name + "/")
    clone2_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone2.variable_scope.name + "/")

    self.assertEqual(net.output_channels, clone1.output_channels)
    self.assertEqual(net.module_name + "_clone", clone1.module_name)
    self.assertEqual("clone2", clone2.module_name)
    self.assertLen(all_vars, 3*len(net_vars))
    self.assertLen(net_vars, len(clone1_vars))
    self.assertLen(net_vars, len(clone2_vars))
    self.assertEqual(net_out.get_shape().as_list(),
                     clone1_out.get_shape().as_list())
    self.assertEqual(net_out.get_shape().as_list(),
                     clone2_out.get_shape().as_list())

  def testDataFormatNotSupported(self):
    """Errors are thrown when an unsupported data_format is used."""

    x = tf.constant(0.0, shape=(2, 8, 6))
    data_format = "WNC"
    self.assertNotIn(data_format, conv.SUPPORTED_1D_DATA_FORMATS)

    with self.assertRaisesRegexp(ValueError, "Invalid data_format"):
      snt.Conv1D(output_channels=4, kernel_shape=4, data_format=data_format)(x)


class Conv1DTransposeTest(parameterized.TestCase, tf.test.TestCase):

  # Constants for use in all tests.
  batch_size = [10, 2, 8, 18, 23]
  in_length = [20, 23, 24, 15, 16]
  in_channels = [6, 2, 3, 6, 9]
  out_channels = [18, 19, 15, 32, 5]
  kernel_shape = [4, 10, 1, 2, 7]
  stride = [1, 2, 4, 7, 5]
  padding = [snt.SAME, snt.SAME, snt.VALID, snt.VALID, snt.VALID]
  use_bias = [True, False, True, False, True]
  out_length = []
  use_output_shape = [True, False, True, False, True]

  for i, pad in enumerate(padding):
    if pad == snt.SAME:
      out_length.append(in_length[i] * stride[i])
    elif pad == snt.VALID:
      out_length.append(in_length[i] * stride[i] + kernel_shape[i] - 1)

  in_shape = tuple(zip(batch_size, in_length, in_channels))
  out_shape = tuple(out_length)
  kernel_shape = tuple(kernel_shape)
  kernel_shape2 = tuple(zip(kernel_shape, out_channels, in_channels))
  stride_shape = tuple(stride)

  def testKernelsNotSpecified(self):
    """Tests error is raised if kernel shape is not specified."""
    with self.assertRaisesRegexp(ValueError, "`kernel_shape` cannot be None."):
      snt.Conv1DTranspose(output_channels=1)

  @parameterized.parameters(
      *zip(out_channels, kernel_shape, padding, use_bias, in_shape, out_shape,
           stride_shape, use_output_shape))
  def testMissingBatchSize(self, out_channels, kernel_shape, padding,
                           use_bias, in_shape, out_shape, stride_shape,
                           use_output_shape):
    """Check functionality with unknown batch size at build time."""

    if use_output_shape:
      output_shape_arg = out_shape
    else:
      output_shape_arg = None
    conv1 = snt.Conv1DTranspose(output_channels=out_channels,
                                output_shape=output_shape_arg,
                                kernel_shape=kernel_shape,
                                padding=padding,
                                stride=stride_shape,
                                name="conv1",
                                use_bias=use_bias)

    # Pass in an image with its batch size set to `None`:
    image = tf.placeholder(tf.float32, shape=(None,) + in_shape[1:])
    output = conv1(image)
    self.assertTrue(output.get_shape().is_compatible_with(
        [None, out_shape, out_channels]))

    with self.test_session() as sess:
      tf.global_variables_initializer().run()
      sess.run(output, feed_dict={image: np.zeros((10,) + in_shape[1:])})

  @parameterized.parameters(
      *zip(batch_size, in_length, in_channels, out_length, out_channels,
           kernel_shape, padding, use_bias, in_shape, out_shape, stride_shape,
           use_output_shape))
  def testShapesSame(self, batch_size, in_length, in_channels, out_length,
                     out_channels, kernel_shape, padding, use_bias, in_shape,
                     out_shape, stride_shape, use_output_shape):
    """The generated shapes are correct."""
    if use_output_shape:
      output_shape_arg = out_shape
    else:
      output_shape_arg = None

    inputs = tf.placeholder(
        tf.float32,
        shape=[batch_size, in_length, in_channels])

    conv1 = snt.Conv1DTranspose(output_channels=out_channels,
                                output_shape=output_shape_arg,
                                kernel_shape=kernel_shape,
                                padding=padding,
                                stride=stride_shape,
                                name="conv1",
                                use_bias=use_bias)

    output = conv1(inputs)

    self.assertTrue(
        output.get_shape().is_compatible_with(
            [batch_size, out_length, out_channels]))

    self.assertTrue(
        conv1.w.get_shape().is_compatible_with(
            [1, kernel_shape, out_channels, in_channels]))

    if use_bias:
      self.assertTrue(
          conv1.b.get_shape().is_compatible_with(
              [out_channels]))

  @parameterized.parameters(
      *zip(out_channels, padding, use_bias, in_shape, out_shape, stride_shape,
           use_output_shape))
  def testKernelShape(self, out_channels, padding, use_bias, in_shape,
                      out_shape, stride_shape, use_output_shape):
    """Errors are thrown for invalid kernel shapes."""
    snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=[3],
        padding=padding,
        stride=stride_shape,
        name="conv1",
        use_bias=use_bias)
    snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=3,
        padding=padding,
        stride=stride_shape,
        name="conv1",
        use_bias=use_bias)

    err = "Invalid kernel"
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1DTranspose(output_channels=out_channels,
                          output_shape=out_shape if use_output_shape else None,
                          kernel_shape=[3, 3],
                          name="conv1",
                          use_bias=use_bias)

    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1DTranspose(output_channels=out_channels,
                          output_shape=out_shape if use_output_shape else None,
                          kernel_shape=[3, 3, 3, 3],
                          name="conv1",
                          use_bias=use_bias)

  @parameterized.parameters(
      *zip(out_channels, padding, use_bias, in_shape, out_shape,
           use_output_shape))
  def testStrideError(self, out_channels, padding, use_bias, in_shape,
                      out_shape, use_output_shape):
    """Errors are thrown for invalid strides."""

    snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=3,
        padding=padding,
        stride=1,
        name="conv1",
        use_bias=use_bias)

    err = ("must be either a positive integer or an iterable of positive "
           "integers of size 1")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1DTranspose(output_channels=out_channels,
                          output_shape=out_shape if use_output_shape else None,
                          kernel_shape=3,
                          padding=padding,
                          stride=[1, 1],
                          name="conv1",
                          use_bias=use_bias)

    with self.assertRaisesRegexp(snt.IncompatibleShapeError, err):
      snt.Conv1DTranspose(output_channels=out_channels,
                          output_shape=out_shape if use_output_shape else None,
                          kernel_shape=3,
                          padding=padding,
                          stride=[1, 1, 1, 1],
                          name="conv1",
                          use_bias=use_bias)

  @parameterized.parameters(
      *zip(batch_size, in_length, in_channels, out_channels, kernel_shape,
           padding, use_bias, out_shape, stride_shape, use_output_shape))
  def testInputTypeError(self, batch_size, in_length, in_channels, out_channels,
                         kernel_shape, padding, use_bias, out_shape,
                         stride_shape, use_output_shape):
    """Errors are thrown for invalid input types."""
    conv1 = snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=kernel_shape,
        padding=padding,
        stride=stride_shape,
        name="conv1",
        use_bias=use_bias)

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([batch_size, in_length,
                               in_channels]), dtype=dtype)
      err = "Input must have dtype tf.float.*"
      with self.assertRaisesRegexp(TypeError, err):
        conv1(x)

  @parameterized.parameters(
      *zip(batch_size, in_length, in_channels, out_channels, kernel_shape,
           padding, use_bias, out_shape, stride_shape, use_output_shape))
  def testSharing(self, batch_size, in_length, in_channels, out_channels,
                  kernel_shape, padding, use_bias, out_shape, stride_shape,
                  use_output_shape):
    """Sharing is working."""

    conv1 = snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=kernel_shape,
        padding=padding,
        stride=stride_shape,
        name="conv1",
        use_bias=use_bias)

    x = np.random.randn(batch_size, in_length, in_channels)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

      # Now change the weights
      w = np.random.randn(1, kernel_shape, out_channels, in_channels)
      conv1.w.assign(w).eval()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

  @parameterized.parameters(
      *zip(batch_size, in_length, in_channels, out_channels, kernel_shape,
           padding, use_bias, out_shape, stride_shape, use_output_shape))
  def testTransposeNWC(self, batch_size, in_length, in_channels, out_channels,
                       kernel_shape, padding, use_bias, out_shape,
                       stride_shape, use_output_shape):
    """Test transpose for NWC format."""

    conv1_transpose = snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=kernel_shape,
        padding=padding,
        stride=stride_shape,
        name="conv1_transpose",
        use_bias=use_bias,
        data_format=conv.DATA_FORMAT_NWC)
    conv1 = conv1_transpose.transpose()

    # Check kernel shapes, strides and padding match.
    self.assertEqual(conv1_transpose.kernel_shape, conv1.kernel_shape)
    self.assertEqual((1, conv1_transpose.stride[1], 1), conv1.stride)
    self.assertEqual(conv1_transpose.conv_op_padding, conv1.conv_op_padding)

    # Before conv1_transpose is connected, we cannot know how many
    # `output_channels` conv1 should have.
    err = "Variables in conv1_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      conv1.output_channels  # pylint: disable=pointless-statement

    # After connection the number of `output_channels` is known.
    x = tf.constant(np.random.randn(batch_size, in_length, in_channels),
                    dtype=np.float32)
    conv1_transpose(x)
    self.assertEqual(in_channels, conv1.output_channels)

    # As is `output_shape`.
    self.assertIn(out_shape, conv1_transpose.output_shape)

    # However, even after connection, the `input_shape` of the forward
    # convolution is not known until it is itself connected (i.e. it can be
    # connected to a different shape input from the `output_shape` of the
    # transpose convolution!)
    err = "Variables in conv1_transpose_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv1.input_shape

  @parameterized.parameters(
      *zip(batch_size, in_length, in_channels, out_channels, kernel_shape,
           padding, use_bias, out_shape, stride_shape, use_output_shape))
  def testTransposeNCW(self, batch_size, in_length, in_channels, out_channels,
                       kernel_shape, padding, use_bias, out_shape,
                       stride_shape, use_output_shape):
    """Test transpose for NCW format."""

    conv1_transpose = snt.Conv1DTranspose(
        output_channels=out_channels,
        output_shape=out_shape if use_output_shape else None,
        kernel_shape=kernel_shape,
        padding=padding,
        stride=stride_shape,
        name="conv1_transpose",
        use_bias=use_bias,
        data_format=conv.DATA_FORMAT_NCW)
    conv1 = conv1_transpose.transpose()

    # Check kernel shapes, strides and padding match.
    self.assertEqual(conv1_transpose.kernel_shape, conv1.kernel_shape)
    self.assertEqual((1, 1, conv1_transpose.stride[2]), conv1.stride)
    self.assertEqual(conv1_transpose.conv_op_padding, conv1.conv_op_padding)

    # Before conv1_transpose is connected, we cannot know how many
    # `output_channels` conv1 should have.
    err = "Variables in conv1_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      conv1.output_channels  # pylint: disable=pointless-statement

    # After connection the number of `output_channels` is known.
    x = tf.constant(np.random.randn(batch_size, in_channels, in_length),
                    dtype=np.float32)
    conv1_transpose(x)
    self.assertEqual(in_channels, conv1.output_channels)

    # As is `output_shape`.
    self.assertIn(out_shape, conv1_transpose.output_shape)

    # However, even after connection, the `input_shape` of the forward
    # convolution is not known until it is itself connected (i.e. it can be
    # connected to a different shape input from the `output_shape` of the
    # transpose convolution!)
    err = "Variables in conv1_transpose_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv1.input_shape

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.Conv1DTranspose(
        output_shape=tf.Dimension(10),
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv1",
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [1, 10, 2]))

    self.assertAllEqual(initializers, initializers_copy)


class CausalConv1DTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputation(self, use_bias):
    """Run through for something with a known answer."""
    conv1 = snt.CausalConv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        use_bias=use_bias,
        name="conv1",
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 1], dtype=np.float32)))
    expected_out = np.reshape(np.array([1, 2, 3, 3, 3]), [1, 5, 1])
    if use_bias:
      expected_out += 1

    init_op = tf.variables_initializer(
        [conv1.w, conv1.b] if use_bias else [conv1.w])
    with self.test_session() as sess:
      sess.run(init_op)
      actual_out = sess.run(out)

    self.assertAllClose(actual_out, expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationStrided(self, use_bias):
    """Run through for something with a known answer."""
    conv1 = snt.CausalConv1D(
        output_channels=1,
        kernel_shape=3,
        stride=2,
        use_bias=use_bias,
        name="conv1",
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 1], dtype=np.float32)))
    expected_out = np.reshape(np.array([1, 3, 3]), [1, 3, 1])
    if use_bias:
      expected_out += 1

    init_op = tf.variables_initializer(
        [conv1.w, conv1.b] if use_bias else [conv1.w])
    with self.test_session() as sess:
      sess.run(init_op)
      actual_out = sess.run(out)

    self.assertAllClose(actual_out, expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationDilated(self, use_bias):
    """Run through for something with a known answer."""
    conv1 = snt.CausalConv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        rate=2,
        use_bias=use_bias,
        name="conv1",
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 1], dtype=np.float32)))
    expected_out = np.reshape(np.array([1, 1, 2, 2, 3]), [1, 5, 1])
    if use_bias:
      expected_out += 1

    init_op = tf.variables_initializer(
        [conv1.w, conv1.b] if use_bias else [conv1.w])
    with self.test_session() as sess:
      sess.run(init_op)
      actual_out = sess.run(out)

    self.assertAllClose(actual_out, expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""

    conv1 = snt.CausalConv1D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        use_bias=use_bias,
        name="conv1")

    x = np.random.randn(1, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    w = np.random.randn(3, 1, 1)
    weight_change_op = conv1.w.assign(w)

    init_op = tf.variables_initializer(
        [conv1.w, conv1.b] if use_bias else [conv1.w])

    with self.test_session() as sess:
      sess.run(init_op)
      first_replica_out = sess.run(out1)
      second_replica_out = sess.run(out2)

      # Now change the weights
      sess.run(weight_change_op)

      first_replica_out_changed = sess.run(out1)
      second_replica_out_changed = sess.run(out2)

    self.assertAllClose(first_replica_out, second_replica_out)
    self.assertAllClose(first_replica_out_changed, second_replica_out_changed)

  def testClone(self):
    net = snt.CausalConv1D(name="conv1d",
                           output_channels=4,
                           kernel_shape=3,
                           stride=5)
    clone1 = net.clone()
    clone2 = net.clone(name="clone2")

    input_to_net = tf.placeholder(tf.float32, shape=[None, 100, 3])
    net_out = net(input_to_net)
    clone1_out = clone1(input_to_net)
    clone2_out = clone2(input_to_net)

    all_vars = tf.trainable_variables()
    net_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=net.variable_scope.name + "/")
    clone1_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone1.variable_scope.name + "/")
    clone2_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone2.variable_scope.name + "/")

    self.assertEqual(net.output_channels, clone1.output_channels)
    self.assertEqual(net.module_name + "_clone", clone1.module_name)
    self.assertEqual("clone2", clone2.module_name)
    self.assertLen(all_vars, 3*len(net_vars))
    self.assertLen(net_vars, len(clone1_vars))
    self.assertLen(net_vars, len(clone2_vars))
    self.assertEqual(net_out.get_shape().as_list(),
                     clone1_out.get_shape().as_list())
    self.assertEqual(net_out.get_shape().as_list(),
                     clone2_out.get_shape().as_list())

  def testDataFormatNotSupported(self):
    """Errors are thrown when an unsupported data_format is used."""

    x = tf.constant(0.0, shape=(2, 8, 6))
    data_format = "WNC"
    self.assertNotIn(data_format, conv.SUPPORTED_1D_DATA_FORMATS)

    with self.assertRaisesRegexp(ValueError, "Invalid data_format"):
      snt.CausalConv1D(output_channels=4, kernel_shape=4,
                       data_format=data_format)(x)


class InPlaneConv2DTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSameNumberOfOutputAndInputChannels(self, use_bias):
    """Test that the number of output and input channels are equal."""

    input_channels = random.randint(1, 32)
    inputs = tf.placeholder(tf.float32, shape=[1, 10, 10, input_channels])
    conv1 = snt.InPlaneConv2D(kernel_shape=3, use_bias=use_bias)

    # Before conv1 is connected, we cannot know how many `output_channels`
    # conv1 should have.
    err = "Variables in in_plane_conv2d not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv1.output_channels

    # After connection, should match `input_channels`.
    conv1(inputs)
    self.assertEqual(conv1.output_channels, input_channels)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""

    conv1 = snt.InPlaneConv2D(kernel_shape=3, use_bias=use_bias)
    x = np.random.randn(1, 5, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)
    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()
      self.assertAllClose(out1.eval(), out2.eval())

      w = np.random.randn(3, 3, 1, 1)  # Now change the weights.
      conv1.w.assign(w).eval()
      self.assertAllClose(out1.eval(), out2.eval())

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.InPlaneConv2D(kernel_shape=3, initializers=initializers)

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))
    self.assertAllEqual(initializers, initializers_copy)


class DepthwiseConv2DTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    """Set up some variables to re-use in multiple tests."""

    super(DepthwiseConv2DTest, self).setUp()

    self.batch_size = batch_size = random.randint(1, 20)
    self.in_height = in_height = random.randint(10, 128)
    self.in_width = in_width = random.randint(10, 128)
    self.in_channels = in_channels = random.randint(1, 10)
    self.kernel_shape_h = kernel_shape_h = random.randint(1, 11)
    self.kernel_shape_w = kernel_shape_w = random.randint(1, 11)
    self.channel_multiplier = channel_multiplier = random.randint(1, 10)
    self.out_channels = out_channels = in_channels * channel_multiplier

    self.input_shape = [batch_size, in_height, in_width, in_channels]
    self.kernel_shape = [kernel_shape_h, kernel_shape_w]
    self.output_shape = [batch_size, in_height, in_width, out_channels]
    self.weight_shape = [kernel_shape_h, kernel_shape_w, in_channels,
                         channel_multiplier]

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesSame(self, use_bias):
    """Test that the generated shapes are correct with SAME padding."""

    out_channels = self.out_channels
    input_shape = self.input_shape
    kernel_shape = self.kernel_shape
    output_shape = self.output_shape
    weight_shape = self.weight_shape
    channel_multiplier = self.channel_multiplier

    inputs = tf.placeholder(tf.float32, shape=input_shape)

    conv1 = snt.DepthwiseConv2D(
        name="conv1",
        channel_multiplier=channel_multiplier,
        kernel_shape=kernel_shape,
        padding=snt.SAME,
        stride=1,
        use_bias=use_bias)
    output = conv1(inputs)

    self.assertEqual(output.get_shape(), output_shape)
    self.assertEqual(conv1.w.get_shape(), weight_shape)
    if use_bias:
      self.assertEqual(conv1.b.get_shape(), out_channels)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesNotKnown(self, use_bias):
    """Test that the generated shapes are correct when input shape not known."""

    inputs = tf.placeholder(
        tf.float32, shape=[None, None, None, self.in_channels], name="inputs")

    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=self.channel_multiplier,
        kernel_shape=self.kernel_shape,
        padding=snt.SAME,
        stride=1,
        use_bias=use_bias)
    output = conv1(inputs)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()
      output_eval = output.eval({inputs: np.zeros(self.input_shape)})
      self.assertEqual(output_eval.shape, tuple(self.output_shape))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testKernelShape(self, use_bias):
    """Test that errors are thrown for invalid kernel shapes."""

    snt.DepthwiseConv2D(channel_multiplier=1, kernel_shape=[3, 4])
    snt.DepthwiseConv2D(channel_multiplier=1, kernel_shape=3)
    error_msg = (r"Invalid kernel shape: x is \[3], must be either a positive"
                 r" integer or an iterable of positive integers of size 2")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, error_msg):
      snt.DepthwiseConv2D(channel_multiplier=1, kernel_shape=[3],
                          use_bias=use_bias, name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testStrideError(self, use_bias):
    """Test that errors are thrown for invalid strides."""

    snt.DepthwiseConv2D(channel_multiplier=1, kernel_shape=3, stride=1,
                        use_bias=use_bias)
    snt.DepthwiseConv2D(channel_multiplier=1, kernel_shape=3, stride=[1] * 2,
                        use_bias=use_bias)
    snt.DepthwiseConv2D(channel_multiplier=1, kernel_shape=3, stride=[1] * 4,
                        use_bias=use_bias)

    error_msg = (r"Invalid stride shape: x is \[1, 1, 1\], must be "
                 r"either a positive integer or an iterable of positive "
                 r"integers of size 2")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, error_msg):
      snt.DepthwiseConv2D(channel_multiplier=3,
                          kernel_shape=3,
                          stride=[1, 1, 1],
                          use_bias=use_bias,
                          name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInputTypeError(self, use_bias):
    """Test that errors are thrown for invalid input types."""
    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=3,
        kernel_shape=3,
        stride=1,
        padding=snt.SAME,
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([1, 5, 5, 1]), dtype=dtype)
      err = "Input must have dtype tf.float.*"
      with self.assertRaisesRegexp(TypeError, err):
        conv1(x)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInitializers(self, use_bias):
    """Test that initializers work as expected."""
    w = random.random()
    b = np.random.randn(6)  # Kernel shape is 3, input channels are 2, 2*3 = 6

    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=3,
        kernel_shape=3,
        stride=1,
        use_bias=use_bias,
        initializers=create_constant_initializers(w, b, use_bias))

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          conv1.w.eval(), np.full(
              [3, 3, 2, 3], w, dtype=np.float32))

      if use_bias:
        self.assertAllClose(conv1.b.eval(), b)

    error_msg = "Initializer for 'w' is not a callable function"
    with self.assertRaisesRegexp(TypeError, error_msg):
      snt.DepthwiseConv2D(
          channel_multiplier=3,
          kernel_shape=3,
          stride=1,
          use_bias=use_bias,
          initializers={"w": tf.ones([])})

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRegularizersInRegularizationLosses(self, use_bias):
    regularizers = create_regularizers(use_bias,
                                       contrib_layers.l1_regularizer(scale=0.5))

    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=3,
        kernel_shape=3,
        stride=1,
        regularizers=regularizers,
        use_bias=use_bias,
        name="conv1")
    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    graph_regularizers = tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES)
    self.assertRegexpMatches(graph_regularizers[0].name, ".*l1_regularizer.*")
    if use_bias:
      self.assertRegexpMatches(graph_regularizers[1].name, ".*l1_regularizer.*")

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=3,
        kernel_shape=3,
        stride=1,
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [10, 10, 1, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationSame(self, use_bias):
    """Run through for something with a known answer using SAME padding."""
    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=1,
        kernel_shape=[3, 3],
        stride=1,
        padding=snt.SAME,
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    expected_out = np.array([[5, 7, 7, 7, 5],
                             [7, 10, 10, 10, 7],
                             [7, 10, 10, 10, 7],
                             [7, 10, 10, 10, 7],
                             [5, 7, 7, 7, 5]])
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(np.reshape(out.eval(), [5, 5]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValid(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""
    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=1,
        kernel_shape=[3, 3],
        stride=1,
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    expected_out = np.array([[10, 10, 10],
                             [10, 10, 10],
                             [10, 10, 10]])
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(np.reshape(out.eval(), [3, 3]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValidMultiChannel(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""
    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=1,
        kernel_shape=[3, 3],
        stride=1,
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 3], dtype=np.float32)))
    expected_out = np.array([[[10] * 3] * 3] * 3)
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          np.reshape(out.eval(), [3, 3, 3]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""
    conv1 = snt.DepthwiseConv2D(
        channel_multiplier=3, kernel_shape=3, stride=1, padding=snt.SAME,
        use_bias=use_bias)

    x = np.random.randn(1, 5, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()
      self.assertAllClose(out1.eval(), out2.eval())

      # Kernel shape was set to 3, which is expandeded to [3, 3, 3].
      # Input channels are 1, output channels := in_channels * multiplier.
      # multiplier is kernel_shape[2] == 3. So weight layout must be:
      # (3, 3, 1, 3).
      w = np.random.randn(3, 3, 1, 3)  # Now change the weights.
      conv1.w.assign(w).eval()
      self.assertAllClose(out1.eval(), out2.eval())


class SeparableConv2DTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    """Set up some variables to re-use in multiple tests."""

    super(SeparableConv2DTest, self).setUp()

    self.batch_size = batch_size = random.randint(1, 100)
    self.in_height = in_height = random.randint(10, 188)
    self.in_width = in_width = random.randint(10, 188)
    self.in_channels = in_channels = random.randint(1, 10)
    self.input_shape = [batch_size, in_height, in_width, in_channels]

    self.kernel_shape_h = kernel_shape_h = random.randint(1, 10)
    self.kernel_shape_w = kernel_shape_w = random.randint(1, 10)
    self.channel_multiplier = channel_multiplier = random.randint(1, 10)
    self.kernel_shape = [kernel_shape_h, kernel_shape_w]

    self.out_channels_dw = out_channels_dw = in_channels * channel_multiplier
    self.output_shape = [batch_size, in_height, in_width, out_channels_dw]
    self.depthwise_filter_shape = [
        kernel_shape_h, kernel_shape_w, in_channels, channel_multiplier
    ]
    self.pointwise_filter_shape = [1, 1, out_channels_dw, out_channels_dw]

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesSame(self, use_bias):
    """Test that the generated shapes are correct with SAME padding."""

    out_channels = self.out_channels_dw
    input_shape = self.input_shape
    kernel_shape = self.kernel_shape
    output_shape = self.output_shape
    depthwise_filter_shape = self.depthwise_filter_shape
    pointwise_filter_shape = self.pointwise_filter_shape
    channel_multiplier = self.channel_multiplier

    inputs = tf.placeholder(tf.float32, shape=input_shape)

    conv1 = snt.SeparableConv2D(
        output_channels=out_channels,
        channel_multiplier=channel_multiplier,
        kernel_shape=kernel_shape,
        padding=snt.SAME,
        use_bias=use_bias)

    output = conv1(inputs)

    self.assertTrue(output.get_shape().is_compatible_with(output_shape))
    self.assertTrue(conv1.w_dw.get_shape().is_compatible_with(
        depthwise_filter_shape))
    self.assertTrue(conv1.w_pw.get_shape().is_compatible_with(
        pointwise_filter_shape))
    if use_bias:
      self.assertTrue(conv1.b.get_shape().is_compatible_with([out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesNotKnown(self, use_bias):
    """Test that the generated shapes are correct when input shape not known."""

    inputs = tf.placeholder(
        tf.float32, shape=[None, None, None, self.in_channels], name="inputs")

    conv1 = snt.SeparableConv2D(
        output_channels=self.out_channels_dw,
        channel_multiplier=1,
        kernel_shape=self.kernel_shape,
        padding=snt.SAME,
        use_bias=use_bias)
    output = conv1(inputs)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      output_eval = output.eval({inputs: np.zeros(self.input_shape)})
      self.assertEqual(output_eval.shape, tuple(self.output_shape))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testKernelShape(self, use_bias):
    """Test that errors are thrown for invalid kernel shapes."""

    # No check against output_channels is done yet (needs input size).
    snt.SeparableConv2D(
        output_channels=1,
        channel_multiplier=2,
        kernel_shape=[3, 4],
        name="conv1",
        use_bias=use_bias)
    snt.SeparableConv2D(
        output_channels=1, channel_multiplier=1, kernel_shape=3, name="conv1")

    error_msg = (r"Invalid kernel shape: x is \[3], must be either a positive"
                 r" integer or an iterable of positive integers of size 2")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, error_msg):
      snt.SeparableConv2D(output_channels=1,
                          channel_multiplier=3,
                          kernel_shape=[3],
                          use_bias=use_bias)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testStrideError(self, use_bias):
    """Test that errors are thrown for invalid strides."""

    snt.SeparableConv2D(
        output_channels=1, channel_multiplier=3, kernel_shape=3, stride=1,
        use_bias=use_bias)
    snt.SeparableConv2D(
        output_channels=1, channel_multiplier=3, kernel_shape=3, stride=[1, 1],
        use_bias=use_bias)
    snt.SeparableConv2D(
        output_channels=1,
        channel_multiplier=3,
        kernel_shape=3,
        stride=[1, 1, 1, 1],
        use_bias=use_bias)

    error_msg = (r"Invalid stride shape: x is \[1, 1, 1\], must be "
                 r"either a positive integer or an iterable of positive "
                 r"integers of size 2")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, error_msg):
      snt.SeparableConv2D(output_channels=1,
                          channel_multiplier=3,
                          kernel_shape=3,
                          stride=[1, 1, 1],
                          name="conv1",
                          use_bias=use_bias)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInputTypeError(self, use_bias):
    """Test that errors are thrown for invalid input types."""
    conv1 = snt.SeparableConv2D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=3,
        padding=snt.SAME,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([1, 5, 5, 1]), dtype=dtype)
      err = "Input must have dtype tf.float.*"
      with self.assertRaisesRegexp(TypeError, err):
        conv1(x)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInitializers(self, use_bias):
    """Test that initializers work as expected."""

    w_dw = random.random()
    w_pw = random.random()
    b = np.random.randn(6)  # Kernel shape is 3, input channels are 2, 2*3 = 6.
    conv1 = snt.SeparableConv2D(
        output_channels=6,
        channel_multiplier=3,
        kernel_shape=3,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            w_dw, w_pw, b, use_bias))

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 2]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(
          conv1.w_dw.eval(), np.full(
              [3, 3, 2, 3], w_dw, dtype=np.float32))
      self.assertAllClose(
          conv1.w_pw.eval(), np.full(
              [1, 1, 6, 6], w_pw, dtype=np.float32))

      if use_bias:
        self.assertAllClose(conv1.b.eval(), b)

    error_msg = "Initializer for 'w_dw' is not a callable function"
    with self.assertRaisesRegexp(TypeError, error_msg):
      snt.SeparableConv2D(
          output_channels=3,
          channel_multiplier=1,
          kernel_shape=3,
          stride=1,
          use_bias=use_bias,
          initializers={"w_dw": tf.ones([])})

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.SeparableConv2D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=3,
        stride=1,
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [10, 10, 1, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRegularizersInRegularizationLosses(self, use_bias):
    regularizers = create_separable_regularizers(
        use_bias, contrib_layers.l1_regularizer(scale=0.5))

    conv1 = snt.SeparableConv2D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=3,
        stride=1,
        regularizers=regularizers,
        use_bias=use_bias,
        name="conv1")
    conv1(tf.placeholder(tf.float32, [10, 10, 1, 2]))

    graph_regularizers = tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES)
    self.assertRegexpMatches(graph_regularizers[0].name, ".*l1_regularizer.*")
    self.assertRegexpMatches(graph_regularizers[1].name, ".*l1_regularizer.*")
    if use_bias:
      self.assertRegexpMatches(graph_regularizers[2].name, ".*l1_regularizer.*")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationSame(self, use_bias):
    """Run through for something with a known answer using SAME padding."""

    conv1 = snt.SeparableConv2D(
        output_channels=1,
        channel_multiplier=1,
        kernel_shape=[3, 3],
        padding=snt.SAME,
        name="conv1",
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    expected_out = np.array([[5, 7, 7, 7, 5],
                             [7, 10, 10, 10, 7],
                             [7, 10, 10, 10, 7],
                             [7, 10, 10, 10, 7],
                             [5, 7, 7, 7, 5]])
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(np.reshape(out.eval(), [5, 5]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationSameNon1Rate(self, use_bias):
    """Same as `testComputationSame`, but have a non-default rate."""

    conv1 = snt.SeparableConv2D(
        output_channels=1,
        channel_multiplier=1,
        kernel_shape=[3, 3],
        padding=snt.SAME,
        name="conv1",
        rate=(2, 3),
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    expected_out = np.array([[5, 5, 3, 5, 5],
                             [5, 5, 3, 5, 5],
                             [7, 7, 4, 7, 7],
                             [5, 5, 3, 5, 5],
                             [5, 5, 3, 5, 5]])
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(np.reshape(out.eval(), [5, 5]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValid(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    conv1 = snt.SeparableConv2D(
        output_channels=1,
        channel_multiplier=1,
        kernel_shape=[3, 3],
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 1], dtype=np.float32)))
    expected_out = np.array([[10, 10, 10],
                             [10, 10, 10],
                             [10, 10, 10]])
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(np.reshape(out.eval(), [3, 3]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValidMultiChannel(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    conv1 = snt.SeparableConv2D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=[3, 3],
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 3], dtype=np.float32)))
    expected_out = np.array([[[28] * 3] * 3] * 3)
    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(np.reshape(out.eval(), [3, 3, 3]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValidChannelMultiplier(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    input_channels = 3
    channel_multiplier = 5
    output_channels = input_channels * channel_multiplier
    conv1 = snt.SeparableConv2D(
        output_channels=output_channels,
        channel_multiplier=channel_multiplier,
        kernel_shape=[3, 3],
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    input_data = np.ones([1, 5, 5, input_channels], dtype=np.float32)
    out = conv1(tf.constant(input_data))
    expected_out = np.ones((3, 3, output_channels)) * 136
    if not use_bias:
      expected_out -= 1

    self.assertTrue(out.get_shape().is_compatible_with([1, 3, 3, output_channels
                                                       ]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(np.reshape(out.eval(), [3, 3, output_channels]),
                          expected_out)
      # Each convolution with weight 1 and size 3x3 results in an output of 9.
      # Pointwise filter is [1, 1, input_channels * channel_multiplier = 15, x].
      # Results in 9 * 15 = 135 + 1 bias = 136 as outputs.

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""
    conv1 = snt.SeparableConv2D(
        output_channels=3, channel_multiplier=3, kernel_shape=3,
        use_bias=use_bias)

    x = np.random.randn(1, 5, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      self.assertAllClose(out1.eval(), out2.eval())

      # Kernel shape was set to 3, which is expandeded to [3, 3, 3].
      # Input channels are 1, output channels := in_channels * multiplier.
      # multiplier is kernel_shape[2] == 3. So weight layout must be:
      # (3, 3, 1, 3).
      w_dw = np.random.randn(3, 3, 1, 3)  # Now change the weights.
      w_pw = np.random.randn(1, 1, 3, 3)  # Now change the weights.
      conv1.w_dw.assign(w_dw).eval()
      conv1.w_pw.assign(w_pw).eval()
      self.assertAllClose(out1.eval(), out2.eval())


class SeparableConv1DTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    """Set up some variables to re-use in multiple tests."""

    super(SeparableConv1DTest, self).setUp()

    self.batch_size = batch_size = random.randint(1, 100)
    self.in_width = in_width = random.randint(10, 188)
    self.in_channels = in_channels = random.randint(1, 10)
    self.input_shape = [batch_size, in_width, in_channels]

    self.kernel_shape_w = kernel_shape_w = random.randint(1, 10)
    self.channel_multiplier = channel_multiplier = random.randint(1, 10)
    self.kernel_shape = [kernel_shape_w]

    self.out_channels_dw = out_channels_dw = in_channels * channel_multiplier
    self.output_shape = [batch_size, in_width, out_channels_dw]
    self.depthwise_filter_shape = [
        1, kernel_shape_w, in_channels, channel_multiplier
    ]
    self.pointwise_filter_shape = [1, 1, out_channels_dw, out_channels_dw]

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesSame(self, use_bias):
    """Test that the generated shapes are correct with SAME padding."""

    out_channels = self.out_channels_dw
    input_shape = self.input_shape
    kernel_shape = self.kernel_shape
    output_shape = self.output_shape
    depthwise_filter_shape = self.depthwise_filter_shape
    pointwise_filter_shape = self.pointwise_filter_shape
    channel_multiplier = self.channel_multiplier

    inputs = tf.placeholder(tf.float32, shape=input_shape)

    conv1 = snt.SeparableConv1D(
        output_channels=out_channels,
        channel_multiplier=channel_multiplier,
        kernel_shape=kernel_shape,
        padding=snt.SAME,
        use_bias=use_bias)

    output = conv1(inputs)

    self.assertTrue(output.get_shape().is_compatible_with(output_shape))
    self.assertTrue(conv1.w_dw.get_shape().is_compatible_with(
        depthwise_filter_shape))
    self.assertTrue(conv1.w_pw.get_shape().is_compatible_with(
        pointwise_filter_shape))
    if use_bias:
      self.assertTrue(conv1.b.get_shape().is_compatible_with([out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesNotKnown(self, use_bias):
    """Test that the generated shapes are correct when input shape not known."""

    inputs = tf.placeholder(
        tf.float32, shape=[None, None, self.in_channels], name="inputs")

    conv1 = snt.SeparableConv1D(
        output_channels=self.out_channels_dw,
        channel_multiplier=1,
        kernel_shape=self.kernel_shape,
        padding=snt.SAME,
        use_bias=use_bias)
    output = conv1(inputs)

    with self.test_session() as session:
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      output_eval = session.run(output, {inputs: np.zeros(self.input_shape)})
    self.assertEqual(output_eval.shape, tuple(self.output_shape))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testKernelShape(self, use_bias):
    """Test that errors are thrown for invalid kernel shapes."""

    # No check against output_channels is done yet (needs input size).
    snt.SeparableConv1D(
        output_channels=1,
        channel_multiplier=2,
        kernel_shape=[3],
        name="conv1",
        use_bias=use_bias)
    snt.SeparableConv1D(
        output_channels=1, channel_multiplier=1, kernel_shape=3, name="conv1")

    error_msg = (r"Invalid kernel shape: x is \[3, 3\], must be either a "
                 r"positive integer or an iterable of positive integers of "
                 r"size 1")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, error_msg):
      snt.SeparableConv1D(output_channels=1,
                          channel_multiplier=3,
                          kernel_shape=[3, 3],
                          use_bias=use_bias)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testStrideError(self, use_bias):
    """Test that errors are thrown for invalid strides."""

    snt.SeparableConv1D(
        output_channels=1, channel_multiplier=3, kernel_shape=3, stride=1,
        use_bias=use_bias)
    snt.SeparableConv1D(
        output_channels=1, channel_multiplier=3, kernel_shape=3, stride=[1],
        use_bias=use_bias)
    snt.SeparableConv1D(
        output_channels=1,
        channel_multiplier=3,
        kernel_shape=3,
        stride=[1, 1, 1],
        use_bias=use_bias)

    error_msg = (r"Invalid stride shape: x is \[1, 1\], must be "
                 r"either a positive integer or an iterable of positive "
                 r"integers of size 1")
    with self.assertRaisesRegexp(snt.IncompatibleShapeError, error_msg):
      snt.SeparableConv1D(output_channels=1,
                          channel_multiplier=3,
                          kernel_shape=3,
                          stride=[1, 1],
                          name="conv1",
                          use_bias=use_bias)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInputTypeError(self, use_bias):
    """Test that errors are thrown for invalid input types."""
    conv1 = snt.SeparableConv1D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=3,
        padding=snt.SAME,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([1, 5, 1]), dtype=dtype)
      err = "Input must have dtype tf.float.*"
      with self.assertRaisesRegexp(TypeError, err):
        conv1(x)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInitializers(self, use_bias):
    """Test that initializers work as expected."""

    w_dw = random.random()
    w_pw = random.random()
    b = np.random.randn(6)  # Kernel shape is 3, input channels are 2, 2*3 = 6.
    conv1 = snt.SeparableConv1D(
        output_channels=6,
        channel_multiplier=3,
        kernel_shape=3,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            w_dw, w_pw, b, use_bias))

    conv1(tf.placeholder(tf.float32, [1, 10, 2]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()

      self.assertAllClose(
          conv1.w_dw.eval(), np.full(
              [1, 3, 2, 3], w_dw, dtype=np.float32))
      self.assertAllClose(
          conv1.w_pw.eval(), np.full(
              [1, 1, 6, 6], w_pw, dtype=np.float32))

      if use_bias:
        self.assertAllClose(conv1.b.eval(), b)

    error_msg = "Initializer for 'w_dw' is not a callable function"
    with self.assertRaisesRegexp(TypeError, error_msg):
      snt.SeparableConv1D(
          output_channels=3,
          channel_multiplier=1,
          kernel_shape=3,
          stride=1,
          use_bias=use_bias,
          initializers={"w_dw": tf.ones([])})

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.SeparableConv1D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=3,
        stride=1,
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [10, 1, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRegularizersInRegularizationLosses(self, use_bias):
    regularizers = create_separable_regularizers(
        use_bias, contrib_layers.l1_regularizer(scale=0.5))

    conv1 = snt.SeparableConv1D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=3,
        stride=1,
        regularizers=regularizers,
        use_bias=use_bias,
        name="conv1")
    conv1(tf.placeholder(tf.float32, [10, 1, 2]))

    graph_regularizers = tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES)
    self.assertRegexpMatches(graph_regularizers[0].name, ".*l1_regularizer.*")
    self.assertRegexpMatches(graph_regularizers[1].name, ".*l1_regularizer.*")
    if use_bias:
      self.assertRegexpMatches(graph_regularizers[2].name, ".*l1_regularizer.*")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationSame(self, use_bias):
    """Run through for something with a known answer using SAME padding."""

    conv1 = snt.SeparableConv1D(
        output_channels=1,
        channel_multiplier=1,
        kernel_shape=[3],
        padding=snt.SAME,
        name="conv1",
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    output = conv1(tf.constant(np.ones([1, 5, 1], dtype=np.float32)))
    expected_out = np.array([[[3], [4], [4], [4], [3]]])
    if not use_bias:
      expected_out -= 1

    with self.test_session() as session:
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      output = session.run(output)
      self.assertAllClose(output, expected_out)

  @parameterized.named_parameters(
      ("WithBiasRateInt", True, 2),
      ("WithBiasRateSeq", True, [2]),
      ("WithoutBiasRateInt", False, 2),
      ("WithoutBiasRateSeq", False, [2]))
  def testComputationSameNon1Rate(self, use_bias, rate):
    """Same as `testComputationSame`, but have a non-default rate."""

    conv1 = snt.SeparableConv1D(
        output_channels=1,
        channel_multiplier=1,
        kernel_shape=[3],
        padding=snt.SAME,
        name="conv1",
        rate=rate,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    output = conv1(tf.constant(np.ones([3, 5, 2], dtype=np.float32)))
    expected_out = np.array([[[5], [5], [7], [5], [5]],
                             [[5], [5], [7], [5], [5]],
                             [[5], [5], [7], [5], [5]]])
    if not use_bias:
      expected_out -= 1

    with self.test_session() as session:
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      output = session.run(output)
      self.assertAllClose(output, expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValid(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    conv1 = snt.SeparableConv1D(
        output_channels=1,
        channel_multiplier=1,
        kernel_shape=[3],
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 1], dtype=np.float32)))
    expected_out = np.array([[[4], [4], [4]]])
    if not use_bias:
      expected_out -= 1

    with self.test_session() as session:
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      out = session.run(out)
      self.assertAllClose(out, expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValidMultiChannel(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    conv1 = snt.SeparableConv1D(
        output_channels=3,
        channel_multiplier=1,
        kernel_shape=[3],
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 3], dtype=np.float32)))
    expected_out = np.array([[[10] * 3] * 3] * 1)
    if not use_bias:
      expected_out -= 1

    with self.test_session() as session:
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      out = session.run(out)
      self.assertAllClose(out, expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValidChannelMultiplier(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    input_channels = 3
    channel_multiplier = 5
    output_channels = input_channels * channel_multiplier
    conv1 = snt.SeparableConv1D(
        output_channels=output_channels,
        channel_multiplier=channel_multiplier,
        kernel_shape=[3],
        padding=snt.VALID,
        use_bias=use_bias,
        initializers=create_separable_constant_initializers(
            1.0, 1.0, 1.0, use_bias))

    input_data = np.ones([1, 5, input_channels], dtype=np.float32)
    out = conv1(tf.constant(input_data))
    expected_out = np.ones((1, 3, output_channels)) * 46
    if not use_bias:
      expected_out -= 1

    self.assertTrue(out.get_shape().is_compatible_with([1, 3, output_channels]))

    with self.test_session() as session:
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      out = session.run(out)
      self.assertAllClose(out, expected_out)
      # Each convolution with weight 1 and size 1x3 results in an output of 3.
      # Pointwise filter is [1, 1, input_channels * channel_multiplier = 15, x].
      # Results in 3 * 15 = 45 + 1 bias = 46 as outputs.

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""
    conv1 = snt.SeparableConv1D(
        output_channels=3, channel_multiplier=3, kernel_shape=3,
        use_bias=use_bias)

    x = np.random.randn(1, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w_dw, conv1.w_pw, conv1.b] if use_bias else
          [conv1.w_dw, conv1.w_pw]).run()
      self.assertAllClose(out1.eval(), out2.eval())

      # Kernel shape was set to 3, which is expandeded to [1, 3].
      # Input channels are 1, output channels := in_channels * multiplier.
      # multiplier is kernel_shape[2] == 3. So weight layout must be:
      # (1, 3, 1, 3).
      w_dw = np.random.randn(1, 3, 1, 3)  # Now change the weights.
      w_pw = np.random.randn(1, 1, 3, 3)  # Now change the weights.
      conv1.w_dw.assign(w_dw).eval()
      conv1.w_pw.assign(w_pw).eval()
      self.assertAllClose(out1.eval(), out2.eval())


class Conv3DTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesSame(self, use_bias):
    """The generated shapes are correct with SAME padding."""

    batch_size = random.randint(1, 100)
    in_depth = random.randint(10, 288)
    in_height = random.randint(10, 288)
    in_width = random.randint(10, 288)
    in_channels = random.randint(1, 10)
    out_channels = random.randint(1, 32)
    kernel_shape_d = random.randint(1, 11)
    kernel_shape_h = random.randint(1, 11)
    kernel_shape_w = random.randint(1, 11)

    inputs = tf.placeholder(
        tf.float32,
        shape=[batch_size, in_depth, in_height, in_width, in_channels])

    conv1 = snt.Conv3D(
        output_channels=out_channels,
        kernel_shape=[kernel_shape_d, kernel_shape_h, kernel_shape_w],
        padding=snt.SAME,
        stride=1,
        use_bias=use_bias,
        name="conv1")

    output = conv1(inputs)

    self.assertTrue(
        output.get_shape().is_compatible_with(
            [batch_size, in_depth, in_height, in_width, out_channels]))

    self.assertTrue(
        conv1.w.get_shape().is_compatible_with(
            [kernel_shape_d, kernel_shape_h, kernel_shape_w, in_channels,
             out_channels]))
    if use_bias:
      self.assertTrue(
          conv1.b.get_shape().is_compatible_with(
              [out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testShapesWithUnknownInputShape(self, use_bias):
    """The generated shapes are correct when input shape not known."""

    batch_size = 5
    in_depth = in_height = in_width = 32
    in_channels = out_channels = 5
    kernel_shape_d = kernel_shape_h = kernel_shape_w = 3

    inputs = tf.placeholder(
        tf.float32,
        shape=[None, None, None, None, in_channels],
        name="inputs")

    conv1 = snt.Conv3D(
        name="conv1",
        output_channels=out_channels,
        kernel_shape=[kernel_shape_d, kernel_shape_h, kernel_shape_w],
        padding=snt.SAME,
        stride=1,
        use_bias=use_bias)

    output = conv1(inputs)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      output_eval = output.eval({
          inputs: np.zeros([batch_size, in_depth, in_height, in_width,
                            in_channels])})

      self.assertEqual(
          output_eval.shape,
          (batch_size, in_depth, in_height, in_width, out_channels))

  def testKernelShape(self):
    """Errors are thrown for invalid kernel shapes."""

    snt.Conv3D(output_channels=10, kernel_shape=[3, 4, 5], name="conv1")
    snt.Conv3D(output_channels=10, kernel_shape=3, name="conv1")

    with self.assertRaisesRegexp(snt.Error, "Invalid kernel shape.*"):
      snt.Conv3D(output_channels=10, kernel_shape=[3, 3], name="conv1")
      snt.Conv3D(output_channels=10, kernel_shape=[3, 3, 3, 3], name="conv1")

  def testStrideError(self):
    """Errors are thrown for invalid strides."""

    snt.Conv3D(
        output_channels=10, kernel_shape=3, stride=1, name="conv1")
    snt.Conv3D(
        output_channels=10, kernel_shape=3, stride=[1, 1, 1], name="conv1")
    snt.Conv3D(
        output_channels=10, kernel_shape=3, stride=[1, 1, 1, 1, 1],
        name="conv1")

    with self.assertRaisesRegexp(snt.Error, "Invalid stride.*"):
      snt.Conv3D(output_channels=10, kernel_shape=3, stride=[1, 1],
                 name="conv1")
      snt.Conv3D(output_channels=10, kernel_shape=3, stride=[1, 1, 1, 1],
                 name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRateError(self, use_bias):
    """Errors are thrown for invalid dilation rates."""

    snt.Conv3D(
        output_channels=10, kernel_shape=3, rate=1, name="conv1",
        use_bias=use_bias)
    snt.Conv3D(
        output_channels=10, kernel_shape=3, rate=2, name="conv1",
        use_bias=use_bias)

    for rate in [0, 0.5, -1]:
      with self.assertRaisesRegexp(snt.IncompatibleShapeError,
                                   "Invalid rate shape*"):
        snt.Conv3D(output_channels=10,
                   kernel_shape=3,
                   rate=rate,
                   name="conv1")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRateAndStrideError(self, use_bias):
    """Errors are thrown for stride > 1 when using atrous convolution."""
    err = "Cannot have stride > 1 with rate > 1"
    with self.assertRaisesRegexp(snt.NotSupportedError, err):
      snt.Conv3D(output_channels=10, kernel_shape=3,
                 stride=2, rate=2, name="conv1", use_bias=use_bias)
    with self.assertRaisesRegexp(snt.NotSupportedError, err):
      snt.Conv3D(output_channels=10, kernel_shape=3,
                 stride=[2, 2, 1], rate=2, name="conv1", use_bias=use_bias)

  def testInputTypeError(self):
    """Errors are thrown for invalid input types."""
    conv1 = snt.Conv3D(output_channels=1,
                       kernel_shape=3,
                       stride=1,
                       padding=snt.SAME,
                       name="conv1",
                       initializers={
                           "w": tf.constant_initializer(1.0),
                           "b": tf.constant_initializer(1.0),
                       })

    for dtype in (tf.uint32, tf.uint64):
      x = tf.constant(np.ones([1, 5, 5, 5, 1]), dtype=dtype)
      self.assertRaisesRegexp(TypeError, "Input must have dtype tf.float.*",
                              conv1, x)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testInitializers(self, use_bias):
    """Test initializers work as expected."""

    w = random.random()
    b = random.random()
    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(w, b, use_bias))

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 10, 2]))

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          conv1.w.eval(),
          np.full([3, 3, 3, 2, 1], w, dtype=np.float32))

      if use_bias:
        self.assertAllClose(
            conv1.b.eval(),
            [b])

    with self.assertRaises(TypeError):
      snt.Conv3D(output_channels=10, kernel_shape=3, stride=1, name="conv1",
                 initializers={"w": tf.ones([])})

  def testInitializerMutation(self):
    """Test that initializers are not mutated."""

    initializers = {"b": tf.constant_initializer(0)}
    initializers_copy = dict(initializers)

    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        name="conv1",
        initializers=initializers)

    conv1(tf.placeholder(tf.float32, [1, 10, 10, 10, 2]))

    self.assertAllEqual(initializers, initializers_copy)

  def testBiasInitializerIsZeroByDefault(self):
    """Test that the default initializer for the bias consists of zeros."""

    conv1 = snt.Conv3D(
        output_channels=5,
        kernel_shape=3,
        stride=1)

    conv1(tf.placeholder(tf.float32, [5, 10, 10, 10, 7]))

    with self.test_session():
      tf.variables_initializer([conv1.w, conv1.b]).run()

      self.assertAllClose(
          conv1.b.eval(),
          np.zeros([5], dtype=np.float32))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testRegularizersInRegularizationLosses(self, use_bias):
    regularizers = create_regularizers(use_bias,
                                       contrib_layers.l1_regularizer(scale=0.5))

    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        regularizers=regularizers,
        use_bias=use_bias,
        name="conv1")
    conv1(tf.placeholder(tf.float32, [1, 10, 10, 10, 2]))

    graph_regularizers = tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES)
    self.assertRegexpMatches(graph_regularizers[0].name, ".*l1_regularizer.*")
    if use_bias:
      self.assertRegexpMatches(graph_regularizers[1].name, ".*l1_regularizer.*")

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationSame(self, use_bias):
    """Run through for something with a known answer using SAME padding."""

    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=snt.SAME,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 5, 1], dtype=np.float32)))
    expected_out = np.asarray([9, 13, 13, 13, 9, 13, 19, 19, 19, 13, 13, 19, 19,
                               19, 13, 13, 19, 19, 19, 13, 9, 13, 13, 13, 9, 13,
                               19, 19, 19, 13, 19, 28, 28, 28, 19, 19, 28, 28,
                               28, 19, 19, 28, 28, 28, 19, 13, 19, 19, 19, 13,
                               13, 19, 19, 19, 13, 19, 28, 28, 28, 19, 19, 28,
                               28, 28, 19, 19, 28, 28, 28, 19, 13, 19, 19, 19,
                               13, 13, 19, 19, 19, 13, 19, 28, 28, 28, 19, 19,
                               28, 28, 28, 19, 19, 28, 28, 28, 19, 13, 19, 19,
                               19, 13, 9, 13, 13, 13, 9, 13, 19, 19, 19, 13, 13,
                               19, 19, 19, 13, 13, 19, 19, 19, 13, 9, 13, 13,
                               13, 9]).reshape((5, 5, 5))

    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          np.reshape(out.eval(), [5, 5, 5]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testComputationValid(self, use_bias):
    """Run through for something with a known answer using snt.VALID padding."""

    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=snt.VALID,
        name="conv1",
        use_bias=use_bias,
        initializers=create_constant_initializers(1.0, 1.0, use_bias))

    out = conv1(tf.constant(np.ones([1, 5, 5, 5, 1], dtype=np.float32)))
    expected_out = np.asarray([28] * 27).reshape((3, 3, 3))

    if not use_bias:
      expected_out -= 1

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          np.reshape(out.eval(), [3, 3, 3]), expected_out)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testSharing(self, use_bias):
    """Sharing is working."""

    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        stride=1,
        padding=snt.SAME,
        use_bias=use_bias,
        name="conv1")

    x = np.random.randn(1, 5, 5, 5, 1)
    x1 = tf.constant(x, dtype=np.float32)
    x2 = tf.constant(x, dtype=np.float32)

    out1 = conv1(x1)
    out2 = conv1(x2)

    with self.test_session():
      tf.variables_initializer(
          [conv1.w, conv1.b] if use_bias else [conv1.w]).run()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

      # Now change the weights
      w = np.random.randn(3, 3, 3, 1, 1)
      conv1.w.assign(w).eval()

      self.assertAllClose(
          out1.eval(),
          out2.eval())

  def testMask1D(self):
    """1D Masks are applied properly."""
    mask = 0.5 * np.ones((3,), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(5, 5, 5, 5, 5))
    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 67.5 * np.ones((5, 3, 3, 3, 1), dtype=np.float32)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask2D(self):
    """2D Masks are applied properly."""
    mask = np.ones((3, 3), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(5, 5, 5, 5, 5))
    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 135 * np.ones((5, 3, 3, 3, 1), dtype=np.float32)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask3D(self):
    """3D Masks are applied properly."""
    mask = np.ones((3, 3, 3), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(5, 5, 5, 5, 5))
    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 135 * np.ones((5, 3, 3, 3, 1), dtype=np.float32)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask4D(self):
    """4D Masks are applied properly."""
    mask = np.ones((3, 3, 3, 5), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(5, 5, 5, 5, 5))
    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 135 * np.ones((5, 3, 3, 3, 1), dtype=np.float32)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMask5D(self):
    """5D Mask are applied properly."""
    mask = np.ones((3, 3, 3, 5, 1), dtype=np.float32)
    inputs = tf.constant(1.0, shape=(5, 5, 5, 5, 5))
    conv1 = snt.Conv3D(
        output_channels=1,
        kernel_shape=3,
        mask=mask,
        padding=snt.VALID,
        use_bias=False,
        initializers=create_constant_initializers(1.0, 0.0, use_bias=False))
    out = conv1(inputs)
    expected_out = 135 * np.ones((5, 3, 3, 3, 1), dtype=np.float32)
    with self.test_session():
      tf.variables_initializer([conv1.w]).run()
      self.assertAllClose(out.eval(), expected_out)

  def testMaskErrorIncompatibleRank1(self):
    """Errors are thrown for incompatible rank 1 mask."""

    np_mask = np.ones((2,), dtype=np.float32)
    x = tf.constant(0.0, shape=(2, 8, 6, 5, 5), dtype=np.float32)

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv3D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank2(self):
    """Errors are thrown for incompatible rank 2 mask."""

    np_mask = np.ones((5, 2), dtype=np.float32)
    x = tf.constant(0.0, shape=(2, 8, 6, 5, 5), dtype=np.float32)

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv3D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank3(self):
    """Errors are thrown for incompatible rank 3 mask."""

    np_mask = np.ones((5, 5, 2), dtype=np.float32)
    x = tf.constant(0.0, shape=(2, 8, 6, 5, 5), dtype=np.float32)

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv3D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank4(self):
    """Errors are thrown for incompatible rank 4 mask."""

    np_mask = np.ones((5, 5, 5, 2), dtype=np.float32)
    x = tf.constant(0.0, shape=(2, 8, 6, 5, 5), dtype=np.float32)

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv3D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIncompatibleRank5(self):
    """Errors are thrown for incompatible rank 5 mask."""

    np_mask = np.ones((5, 5, 5, 5, 2), dtype=np.float32)
    x = tf.constant(0.0, shape=(2, 8, 6, 5, 5), dtype=np.float32)

    # Test with both numpy arrays and Tensors.
    for mask in (np_mask, tf.convert_to_tensor(np_mask)):
      with self.assertRaises(snt.Error) as cm:
        snt.Conv3D(output_channels=4, kernel_shape=5, mask=mask)(x)
      self.assertTrue(str(cm.exception).startswith(
          "Invalid mask shape: {}".format(np_mask.shape)))

  def testMaskErrorIfIncorrectDtype(self):
    """Errors are thrown when a Tensor with incorrect dtype is used."""

    mask = tf.constant(0, shape=(4, 4), dtype=tf.int32)
    x = tf.constant(0.0, shape=(2, 8, 6, 5, 5))

    with self.assertRaises(TypeError) as cm:
      snt.Conv3D(output_channels=4, kernel_shape=(4, 4, 4), mask=mask)(x)
    self.assertTrue(str(cm.exception).startswith(
        "Mask needs to have dtype float16, bfloat16, float32 or float64"))

  def testClone(self):
    net = snt.Conv3D(name="conv3d",
                     output_channels=4,
                     kernel_shape=3,
                     stride=5)
    clone1 = net.clone()
    clone2 = net.clone(name="clone2")

    input_to_net = tf.placeholder(tf.float32, shape=[None, 101, 102, 103, 3])
    net_out = net(input_to_net)
    clone1_out = clone1(input_to_net)
    clone2_out = clone2(input_to_net)

    all_vars = tf.trainable_variables()
    net_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=net.variable_scope.name + "/")
    clone1_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone1.variable_scope.name + "/")
    clone2_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=clone2.variable_scope.name + "/")

    self.assertEqual(net.output_channels, clone1.output_channels)
    self.assertEqual(net.module_name + "_clone", clone1.module_name)
    self.assertEqual("clone2", clone2.module_name)
    self.assertLen(all_vars, 3*len(net_vars))
    self.assertLen(net_vars, len(clone1_vars))
    self.assertLen(net_vars, len(clone2_vars))
    self.assertEqual(net_out.get_shape().as_list(),
                     clone1_out.get_shape().as_list())
    self.assertEqual(net_out.get_shape().as_list(),
                     clone2_out.get_shape().as_list())

  def testDataFormatNotSupported(self):
    """Errors are thrown when an unsupported data_format is used."""

    x = tf.constant(0.0, shape=(2, 7, 8, 9, 6))
    data_format = "NCHWD"
    self.assertNotIn(data_format, conv.SUPPORTED_3D_DATA_FORMATS)

    with self.assertRaisesRegexp(ValueError, "Invalid data_format"):
      snt.Conv3D(output_channels=4, kernel_shape=4, data_format=data_format)(x)


class Conv3DTransposeTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    """Set up some variables to re-use in multiple tests."""

    super(Conv3DTransposeTest, self).setUp()

    self.batch_size = 7
    self.in_depth = 7
    self.in_height = 7
    self.in_width = 11
    self.in_channels = 4
    self.out_channels = 10
    self.kernel_shape_d = 5
    self.kernel_shape_h = 5
    self.kernel_shape_w = 7
    self.stride_d = 1
    self.stride_h = 2
    self.stride_w = 3
    self.padding = snt.SAME

    self.in_shape = (self.batch_size, self.in_depth, self.in_height,
                     self.in_width, self.in_channels)

    self.out_shape = (self.in_depth, self.in_height, self.in_width)

    self.kernel_shape = (self.kernel_shape_d, self.kernel_shape_h,
                         self.kernel_shape_w)

    self.kernel_shape2 = (self.kernel_shape_d, self.kernel_shape_h,
                          self.kernel_shape_w, self.out_channels,
                          self.in_channels)

    self.strides = (self.stride_d, self.stride_h, self.stride_w)

  def testKernelsNotSpecified(self):
    with self.assertRaisesRegexp(ValueError, "`kernel_shape` cannot be None."):
      snt.Conv3DTranspose(output_channels=1)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testOutputShapeConsistency(self, use_bias):
    """Tests if output shapes are valid."""

    # When padding is SAME, then the actual number of padding pixels can be
    # computed as: pad = kernel_shape - strides + (-input_shape % strides)
    #                 =     5         -    1    + (- 32       %      1) = 4

    # The formula for the minimal size is:
    # oH = strides[1] * (in_height - 1) - padding + kernel_shape_h
    # oH =          1 * (       32 - 1) -    4    +       5 = 32

    # The formula for the maximum size (due to extra pixels) is:
    # oH_max = oH + strides[1] - 1
    # so, for strides = 1 and padding = SAME, input size == output size.
    inputs = tf.placeholder(tf.float32, shape=self.in_shape)

    conv1 = snt.Conv3DTranspose(name="conv3d_1",
                                output_channels=self.out_channels,
                                output_shape=self.out_shape,
                                kernel_shape=self.kernel_shape,
                                padding=self.padding,
                                stride=1,
                                use_bias=use_bias)

    outputs = conv1(inputs)

    self.assertTrue(outputs.get_shape().is_compatible_with((
        self.batch_size,) + self.out_shape + (self.out_channels,)))

    self.assertTrue(conv1.w.get_shape().is_compatible_with(self.kernel_shape2))
    if use_bias:
      self.assertTrue(conv1.b.get_shape().is_compatible_with(
          [self.out_channels]))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testOutputShapeInteger(self, use_bias):
    """Tests if output shapes are valid when specified as an integer."""
    inputs = tf.zeros(shape=[3, 5, 5, 5, 2], dtype=tf.float32)
    inputs_2 = tf.zeros(shape=[3, 5, 7, 5, 2], dtype=tf.float32)

    conv1 = snt.Conv3DTranspose(name="conv3d_1",
                                output_channels=10,
                                output_shape=10,
                                kernel_shape=5,
                                padding=snt.SAME,
                                stride=2,
                                use_bias=use_bias)

    outputs = conv1(inputs)
    outputs_2 = conv1(inputs_2)

    self.assertTrue(outputs.get_shape().is_compatible_with((3, 10, 10, 10, 10)))

    with self.test_session() as sess:
      tf.global_variables_initializer().run()
      sess.run(outputs)
      with self.assertRaises(tf.errors.InvalidArgumentError):
        sess.run(outputs_2)

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testOutputShapeInference(self, use_bias):
    """Tests if output shapes are valid when not specified."""
    inputs = tf.zeros(shape=[3, 5, 5, 5, 2], dtype=tf.float32)

    conv1 = snt.Conv3DTranspose(name="conv3d_1",
                                output_channels=10,
                                output_shape=None,
                                kernel_shape=5,
                                padding=snt.SAME,
                                stride=2,
                                use_bias=use_bias)

    outputs = conv1(inputs)

    self.assertTrue(outputs.get_shape().is_compatible_with((3, 10, 10, 10, 10)))

  def testBiasInitializerIsZeroByDefault(self):
    """Test that the default initializer for the bias consists of zeros."""

    conv1 = snt.Conv3DTranspose(
        output_channels=7,
        kernel_shape=3,
        stride=1)

    conv1(tf.placeholder(tf.float32, [7, 10, 10, 10, 5]))

    with self.test_session():
      tf.variables_initializer([conv1.w, conv1.b]).run()

      self.assertAllClose(
          conv1.b.eval(),
          np.zeros([7], dtype=np.float32))

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testTransposition(self, use_bias):
    """Tests if the correct output shapes are setup in transposed module."""
    net = snt.Conv3DTranspose(name="conv3d_3",
                              output_channels=self.out_channels,
                              output_shape=self.out_shape,
                              kernel_shape=self.kernel_shape,
                              padding=self.padding,
                              stride=1,
                              use_bias=use_bias)

    net_transpose = net.transpose()
    input_to_net = tf.placeholder(tf.float32, shape=self.in_shape)
    err = "Variables in {} not instantiated yet, __call__ the module first."
    with self.assertRaisesRegexp(snt.NotConnectedError,
                                 err.format(net.scope_name)):
      net_transpose(input_to_net)
    net_transpose = net.transpose(name="another_net_transpose")
    net_out = net(input_to_net)
    net_transposed_output = net_transpose(net_out)
    self.assertEqual(net_transposed_output.get_shape(),
                     input_to_net.get_shape())

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testTransposeNDHWC(self, use_bias):
    """Test transpose for NDHWC format."""

    conv3_transpose = snt.Conv3DTranspose(
        output_channels=self.out_channels,
        output_shape=self.out_shape,
        kernel_shape=self.kernel_shape,
        padding=self.padding,
        stride=self.strides,
        name="conv3_transpose",
        use_bias=use_bias,
        data_format=conv.DATA_FORMAT_NDHWC)
    conv3 = conv3_transpose.transpose()

    # Check kernel shapes, strides and padding match.
    self.assertEqual(conv3_transpose.kernel_shape, conv3.kernel_shape)
    self.assertEqual((1,) + self.strides + (1,), conv3.stride)
    self.assertEqual(conv3_transpose.conv_op_padding, conv3.conv_op_padding)

    # Before conv3_transpose is connected, we cannot know how many
    # `output_channels` conv1 should have.
    err = "Variables in conv3_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv3.output_channels

    # After connection the number of `output_channels` is known.
    x = tf.constant(np.random.randn(self.batch_size,
                                    self.in_depth,
                                    self.in_height,
                                    self.in_width,
                                    self.in_channels),
                    dtype=np.float32)
    conv3_transpose(x)
    self.assertEqual(self.in_channels, conv3.output_channels)

    # However, even after connection, the `input_shape` of the forward
    # convolution is not known until it is itself connected (i.e. it can be
    # connected to a different shape input from the `output_shape` of the
    # transpose convolution!)
    err = "Variables in conv3_transpose_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv3.input_shape

  @parameterized.named_parameters(
      ("WithBias", True),
      ("WithoutBias", False))
  def testTransposeNCDHW(self, use_bias):
    """Test transpose for NCDHW format."""

    conv3_transpose = snt.Conv3DTranspose(
        output_channels=self.out_channels,
        output_shape=self.out_shape,
        kernel_shape=self.kernel_shape,
        padding=self.padding,
        stride=self.strides,
        name="conv3_transpose",
        use_bias=use_bias,
        data_format=conv.DATA_FORMAT_NCDHW)
    conv3 = conv3_transpose.transpose()

    # Check kernel shapes, strides and padding match.
    self.assertEqual(conv3_transpose.kernel_shape, conv3.kernel_shape)
    self.assertEqual((1, 1) + self.strides, conv3.stride)
    self.assertEqual(conv3_transpose.conv_op_padding, conv3.conv_op_padding)

    # Before conv3_transpose is connected, we cannot know how many
    # `output_channels` conv1 should have.
    err = "Variables in conv3_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv3.output_channels

    # After connection the number of `output_channels` is known.
    x = tf.constant(np.random.randn(self.batch_size,
                                    self.in_channels,
                                    self.in_depth,
                                    self.in_height,
                                    self.in_width),
                    dtype=np.float32)
    conv3_transpose(x)
    self.assertEqual(self.in_channels, conv3.output_channels)

    # However, even after connection, the `input_shape` of the forward
    # convolution is not known until it is itself connected (i.e. it can be
    # connected to a different shape input from the `output_shape` of the
    # transpose convolution!)
    err = "Variables in conv3_transpose_transpose not instantiated yet"
    with self.assertRaisesRegexp(snt.NotConnectedError, err):
      _ = conv3.input_shape

if __name__ == "__main__":
  tf.test.main()
