import tensorflow as tf

from tensorflow.compiler.mlir.tfr.python import test_utils
from tensorflow.python.framework import test_ops
from tensorflow.python.platform import test


class TestOpsDefsTest(test_utils.OpsDefsTest):

  def test_test_ops(self):
    attr = tf.function(test_ops.test_attr)(tf.float32)
    self.assertAllClose(attr.numpy(), 100.0)


if __name__ == '__main__':
  os.environ['TF_MLIR_TFR_LIB_DIR'] = (

"""Test utils for composite op definition."""
from tensorflow.python.eager import backprop
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test


class OpsDefsTest(test.TestCase):
  """Test utils."""

  def _assertOpAndComposite(self, vars_, compute_op, compute_composite, kwargs,
                            op_kwargs=None):
    if op_kwargs is None:
      op_kwargs = kwargs

_lib_dir = os.path.dirname(gen_pad_ops.__file__)
_lib_name = os.path.basename(gen_pad_ops.__file__)[4:].replace('.py', '.so')
load_library.load_op_library(os.path.join(_lib_dir, _lib_name))


class PadOpsDefsTest(test_utils.OpsDefsTest, parameterized.TestCase):

  @parameterized.named_parameters(('ReflectMode', 'REFLECT'),
                                  ('SymmetricMode', 'SYMMETRIC'))
  def test_mirror_pad(self, mode):

_lib_dir = os.path.dirname(gen_mnist_ops.__file__)
_lib_name = os.path.basename(gen_mnist_ops.__file__)[4:].replace('.py', '.so')
load_library.load_op_library(os.path.join(_lib_dir, _lib_name))


class MnistOpsDefsTest(test_utils.OpsDefsTest):

  def test_new_conv2d_relu(self):
    input_ = tf.random.uniform([1, 4, 4, 1])
    filter_ = tf.random.uniform([2, 2, 1, 8])
    bias = tf.zeros([8])
    kwargs = {
        'input_': input_,