from tensorflow.python.util.lazy_loader import KerasLazyLoader as _KerasLazyLoader

# Make sure code inside the TensorFlow codebase can use tf2.enabled() at import.
_os.environ["TF2_BEHAVIOR"] = "1"
from tensorflow.python import tf2 as _tf2
_tf2.enable()

# API IMPORTS PLACEHOLDER

# WRAPPER_PLACEHOLDER

# Make sure directory containing top level submodules is in

}

void CreateReplicatedClusteringPipeline(OpPassManager &pm,
                                        llvm::StringRef module_name) {
  // Since the internal bridge clustering passes are shared among TF1/TF2
  // TF2-only passes should go here. However, this should be very rare and
  // new passes generally should go into the internal
  // AddReplicatedBridgeClusteringPipelinePasses.

// enabled or disabled, it will decide whether or not to run the bridge.
//
// The config_proto param is a required input for all TF1 graphs but it is
// redundant for TF2 graphs.
// If getting rollout policy involves graph analysis, `record_stats` is used
// to decide whether to emit metrics on unsupported features of the graph.
MlirBridgeRolloutPolicy GetMlirBridgeRolloutPolicy(

        "//tensorflow/cc/experimental/libtf:__pkg__",
        "//tensorflow/cc/experimental/libtf:__subpackages__",
        # copybara:uncomment_begin(google-only)
        # "//tensorflow/cc/experimental/tf2:__pkg__",
        # "//tensorflow/cc/experimental/tf2:__subpackages__",
        # copybara:uncomment_end
        "//tensorflow/compiler/mlir/tensorflow/c:__subpackages__",
        "//tensorflow/core/transforms:__subpackages__",
    ],

limitations under the License.
==============================================================================*/

// This pass separates SparseCore, TensorCore, and non-TPU operations into
// separate functions for proper sequencing of TF2 TPU Embedding (see
// tpu_embedding_v3.py). This pass is a precursor for pipelining (see
// embedding_pipelining.cc) and DOES NOT permit parallel execution across SC and

_CalibrationMethod = qc.CalibrationOptions.CalibrationMethod


# Test cases for Static Range Quantization.
# Tries to run all tests cases in both the graph mode (default in TF1) and the
# eager mode (default in TF2) to ensure support for when TF2 is disabled.
class StaticRangeQuantizationTest(quantize_model_test_base.QuantizedModelTest):

  @parameterized.parameters(
      testing.parameter_combinations([{
          'bias_fn': (
              None,

  return *jitrt_flags;
}

ConfigProto::Experimental::MlirBridgeRollout GetMlirBridgeRolloutState(
    std::optional<const ConfigProto> config_proto) {
  // TF1 graphs that do not override Sessions's ConfigProto and TF2 graphs
  // can enable/disable the graph via tf_mlir_enable_mlir_bridge.
  auto tf_mlir_enable_mlir_bridge =
      GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge;
  if (tf_mlir_enable_mlir_bridge !=

                          status.ToString())
      .IgnoreError();

  return status;
}

// V1 Compat Bridge takes a TF Executor dialect and extracts the TF2 portion
// and inserts it into a submodule. We just want to run the clustering
// portion of the pipeline on just the single submodule.
absl::Status RunClusteringPipelineOnSubmodule(

    representative_dataset: rd.RepresentativeDataset,
) -> None:
  """Runs the representative dataset through a function for calibration.

  NOTE: This is intended to be run in eager mode (TF2).

  Args:
    func: The function to run the representative samples through.
    representative_dataset: Representative dataset used for calibration. The

        "//tensorflow/core:tensorflow",
        "//tensorflow/core:test",
        "//tensorflow/core:test_main",
        "//tensorflow/core:testlib",
    ],
)

# A subset of the TF2 saved models can be generated with this tool.
py_strict_binary(
    name = "testdata/generate_saved_models",
    srcs = ["testdata/generate_saved_models.py"],
    data = [
        ":saved_model_asset_data",