# Path: mesh_tensorflow/test_utils.py
# class NumpyConverter(object):
#   def __init__(self):
#   def convert_np_array_to_mtf_tensor(self, x, dim_names=None, dtype=tf.int32):
#   def convert_mtf_tensor_to_np_array(self, x_mtf):
#   def convert_mtf_tensor_to_tf_tensor(self, mtf_tensor):
#   def graph(self):
#   def mesh(self):
#   def session(self):


from absl.testing import parameterized
from mesh_tensorflow import test_utils
from tensorflow.python.framework import test_util  # pylint:disable=g-direct-tensorflow-import
import mesh_tensorflow as mtf
import mock
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_probability as tfp
    length = 4
    filter_size = 3

    x = np.random.randn(batch, length, d_model)
    x_mtf = converter.convert_np_array_to_mtf_tensor(
        x, dtype=tf.float32, dim_names=["batch", "length", "d_model"])

    conv_filter = np.random.randn(1, filter_size, d_model, d_out)
    initializer = lambda shape, dtype, **kwargs: conv_filter
    output_mtf = mtf.layers.conv1d(
        x_mtf,
        output_dim=mtf.Dimension("output_dim", d_out),
        filter_size=filter_size,
        padding="VALID",
        filter_initializer=initializer)
    actual = converter.convert_mtf_tensor_to_np_array(output_mtf)

    # Expected length is 2.
    expected = np.empty(shape=(batch, 2, d_out), dtype=np.float32)

    # [filter_size, d_model]
    current_filter = conv_filter[0, :, :, 0]
    # b: batch, k: filter_size, d: d_model.
    expected[:, 0] = np.einsum("bkd,kd->b", x[:, :filter_size, :],
                               current_filter).reshape(batch, 1)
    expected[:, 1] = np.einsum("bkd,kd->b", x[:, 1:, :],
                               current_filter).reshape(batch, 1)
    self.assertAllClose(actual, expected)

  def testConv1dValidPaddingMultipleBatchDims(self):

# Path: mesh_tensorflow/test_utils.py
# class NumpyConverter(object):
#   def __init__(self):
#   def convert_np_array_to_mtf_tensor(self, x, dim_names=None, dtype=tf.int32):
#   def convert_mtf_tensor_to_np_array(self, x_mtf):
#   def convert_mtf_tensor_to_tf_tensor(self, mtf_tensor):
#   def graph(self):
#   def mesh(self):
#   def session(self):


# Path: mesh_tensorflow/layers_test.py
from absl.testing import parameterized
from mesh_tensorflow import test_utils
from tensorflow.python.framework import test_util  # pylint:disable=g-direct-tensorflow-import
import mesh_tensorflow as mtf
import mock
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_probability as tfp

    length = 4
    filter_size = 3

    x = np.random.randn(batch, length, d_model)
    x_mtf = converter.convert_np_array_to_mtf_tensor(
        x, dtype=tf.float32, dim_names=["batch", "length", "d_model"])

    conv_filter = np.random.randn(1, filter_size, d_model, d_out)
    initializer = lambda shape, dtype, **kwargs: conv_filter
    output_mtf = mtf.layers.conv1d(
        x_mtf,
        output_dim=mtf.Dimension("output_dim", d_out),
        filter_size=filter_size,
        padding="VALID",
        filter_initializer=initializer)
    actual = converter.convert_mtf_tensor_to_np_array(output_mtf)

    # Expected length is 2.
    expected = np.empty(shape=(batch, 2, d_out), dtype=np.float32)

    # [filter_size, d_model]
    current_filter = conv_filter[0, :, :, 0]
    # b: batch, k: filter_size, d: d_model.
    expected[:, 0] = np.einsum("bkd,kd->b", x[:, :filter_size, :],
                               current_filter).reshape(batch, 1)
    expected[:, 1] = np.einsum("bkd,kd->b", x[:, 1:, :],
                               current_filter).reshape(batch, 1)
    self.assertAllClose(actual, expected)

  def testConv1dValidPaddingMultipleBatchDims(self):

    converter = test_utils.NumpyConverter()