# 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 as mtf_test_utils
from tensorflow.python.framework import test_util  # pylint:disable=g-direct-tensorflow-import
import mesh_tensorflow as mtf
import numpy as np
import tensorflow.compat.v1 as tf
  })
  def testPoolTensor1d(self, pool_fn, pool_fn_mtf):
    converter = mtf_test_utils.NumpyConverter()
    pool_size = 2
    x = np.random.randn(2, 3, 4, 5)
    expected = np.empty(shape=[2, 3, 2, 5])
    expected[:, :, 0, :] = pool_fn(x[:, :, 0:2, :], axis=2)
    expected[:, :, 1, :] = pool_fn(x[:, :, 2:4, :], axis=2)

    x_mtf = converter.convert_np_array_to_mtf_tensor(x, dtype=tf.float32)
    pooled_mtf = mtf.pool_tensor_1d(
        x_mtf,
        pool_dim=x_mtf.shape.dims[2],
        reduce_fn=pool_fn_mtf,
        pool_size=pool_size)
    actual = converter.convert_mtf_tensor_to_np_array(pooled_mtf)
    self.assertAllClose(expected, actual)

  @parameterized.parameters({"pool_size": 2}, {"pool_size": 3})
  def testStrideTensor1d(self, pool_size):
    converter = mtf_test_utils.NumpyConverter()
    x = np.random.randint(0, 100, size=[2, 3, 6, 5])
    x_mtf = converter.convert_np_array_to_mtf_tensor(x)
    expected = x[:, :, range(0, x.shape[2], pool_size), :]
    strided_mtf = mtf.stride_tensor_1d(
        x_mtf, pool_dim=x_mtf.shape.dims[2], pool_size=pool_size)
    actual = converter.convert_mtf_tensor_to_np_array(strided_mtf)
    self.assertAllEqual(expected, actual)

  def testReduceFirst(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/ops_test.py
from absl.testing import parameterized
from mesh_tensorflow import test_utils as mtf_test_utils
from tensorflow.python.framework import test_util  # pylint:disable=g-direct-tensorflow-import
import mesh_tensorflow as mtf
import numpy as np
import tensorflow.compat.v1 as tf

  })
  def testPoolTensor1d(self, pool_fn, pool_fn_mtf):
    converter = mtf_test_utils.NumpyConverter()
    pool_size = 2
    x = np.random.randn(2, 3, 4, 5)
    expected = np.empty(shape=[2, 3, 2, 5])
    expected[:, :, 0, :] = pool_fn(x[:, :, 0:2, :], axis=2)
    expected[:, :, 1, :] = pool_fn(x[:, :, 2:4, :], axis=2)

    x_mtf = converter.convert_np_array_to_mtf_tensor(x, dtype=tf.float32)
    pooled_mtf = mtf.pool_tensor_1d(
        x_mtf,
        pool_dim=x_mtf.shape.dims[2],
        reduce_fn=pool_fn_mtf,
        pool_size=pool_size)
    actual = converter.convert_mtf_tensor_to_np_array(pooled_mtf)
    self.assertAllClose(expected, actual)

  @parameterized.parameters({"pool_size": 2}, {"pool_size": 3})
  def testStrideTensor1d(self, pool_size):
    converter = mtf_test_utils.NumpyConverter()
    x = np.random.randint(0, 100, size=[2, 3, 6, 5])
    x_mtf = converter.convert_np_array_to_mtf_tensor(x)
    expected = x[:, :, range(0, x.shape[2], pool_size), :]
    strided_mtf = mtf.stride_tensor_1d(
        x_mtf, pool_dim=x_mtf.shape.dims[2], pool_size=pool_size)
    actual = converter.convert_mtf_tensor_to_np_array(strided_mtf)
    self.assertAllEqual(expected, actual)

  def testReduceFirst(self):

    converter = mtf_test_utils.NumpyConverter()