from tensorflow.python.module import module
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import saved_model

TESTDATA_PATH = flags.DEFINE_string(
    "path", None, help="Path to testdata directory.")

MODEL_NAME = flags.DEFINE_string(
    "model_name", None, help="Name of model to generate.")


class DataStructureModel(module.Module):
  """Model used for testing data structures in the C++ API."""

from absl import app
from absl import flags

# TODO(b/263048929): Create a test for gen_quantized_function_library.

_OUTPUT_FILE = flags.DEFINE_string('output_file', None, 'output file location')
_SRCS = flags.DEFINE_string('src', None, 'source file locations')
_NAMESPACE = flags.DEFINE_string('namespace', 'mlir::quant',
                                 'namespace in the generated file')

flags.mark_flags_as_required(['output_file', 'src'])

from tensorflow.compiler.mlir.tfr.python.tfr_gen import tfr_gen_from_module
from tensorflow.python.platform import flags

FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output', None,
    'Path to write the genereated register op file and MLIR file.')

flags.DEFINE_bool('gen_register_op', True,
                  'Generate register op cc file or tfr mlir file.')

from tensorflow.python.framework import dtypes
from tensorflow.python.ops import gen_array_ops as array_ops
from tensorflow.python.platform import flags


Composite = composite.Composite
FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output', None,
    'Path to write the genereated register op file and MLIR file.')

flags.DEFINE_bool('gen_register_op', True,
                  'Generate register op cc file or tfr mlir file.')

from tensorflow.python.util import compat
from tensorflow.tools.proto_splitter import constants
from tensorflow.tools.proto_splitter.python import saved_model as proto_splitter

SPLITTER_TESTDATA_PATH = flags.DEFINE_string(
    "path", None, help="Path to testdata directory.")


def generate_non_chunked_model(non_chunked_dir: str):
  root = module.Module()
  root.c = constant_op.constant(np.random.random_sample([150, 150]))

from tensorflow.python import pywrap_mlir  # pylint: disable=g-direct-tensorflow-import
from tensorflow.python.lib.io import file_io

# Use /tmp to make debugging the tests easier (see README.md)
flags.DEFINE_string('save_model_path', '', 'Path to save the model to.')
FLAGS = flags.FLAGS


# This function needs to take a "create_module_fn", as opposed to just the

from absl import logging
import tensorflow.compat.v1 as tf

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

# Use /tmp to make debugging the tests easier (see README.md)
flags.DEFINE_string('save_model_path', '', 'Path to save the model to.')
FLAGS = flags.FLAGS


def set_tf_options():
  # Default TF1.x uses reference variables that are not supported by SavedModel

from tensorflow.python.ops import gen_array_ops
from tensorflow.python.platform import flags

Composite = composite.Composite
FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output', None,
    'Path to write the genereated register op file and MLIR file.')

flags.DEFINE_bool('gen_register_op', True,
                  'Generate register op cc file or tfr mlir file.')

from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import flags


Composite = composite.Composite
FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output', None,
    'Path to write the genereated register op file and MLIR file.')

flags.DEFINE_bool('gen_register_op', True,
                  'Generate register op cc file or tfr mlir file.')