Raises:
      KeyError: If the set of input keys in the dataset samples doesn't match
      the set of expected input keys.
    """
    # When running in graph mode (TF1), tf.Tensor types should be converted to
    # numpy ndarray types to be compatible with `make_tensor_proto`.
    if not context.executing_eagerly():
      with session.Session() as sess:

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

  NOTE: This is intended to be run in graph mode (TF1).

  The function is identified by the SignatureDef.

  Args:
    sess: The Session object to run the function in.
    signature_def: A SignatureDef that identifies a function by specifying the

  absl::call_once(flags_init, &AllocateAndParseFlags);
  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;

      filter_shape: Sequence[int] = (2, 3, 3, 2),
      use_variable_for_filter=False,
  ) -> Tuple[core.Tensor, core.Tensor]:
    """Creates a basic convolution model.

    This is intended to be used for TF1 (graph mode) tests.

    Args:
      input_shape: Shape of the input tensor.
      filter_shape: Shape of the filter.
      use_variable_for_filter: Setting this to `True` makes the filter for the

  return status;
}

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.

  // or ambiguity in this unit-wise precision, our quantizer will raise an
  // error.
  repeated UnitWiseQuantizationSpec unit_wise_quantization_specs = 17;

  // (TF1 SavedModel only) Collection of tags identifying the MetaGraphDef
  // within the SavedModel to analyze. If not specified, ["serve"] is used.
  repeated string tags = 5;

    const std::string& directory,
    const absl::optional<std::unordered_set<std::string>>& tags,
    ImmediateExecutionContext* context, std::unique_ptr<TFSavedModelAPI>* out) {
  // TODO(bmzhao): Add support for loading a TF1 SavedModel.
  if (tags) {
    return errors::Unimplemented(
        "Loading saved models with explicit tags will be supported in the "
        "future");
  }

  SavedModelV2Bundle bundle;

from tensorflow.python.types import core

_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(

@test_util.run_all_in_graph_and_eager_modes
class QuantizationOptionsTest(quantize_model_test_base.QuantizedModelTest):
  """Test cases regarding the use of QuantizationOptions proto.

  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 SimpleModel(module.Module):

    def __init__(self):

        checkpoint, please change your optimizer to
        `tf.keras.optimizers.legacy.XXX` (e.g.
        `tf.keras.optimizers.legacy.Adam`).
    *   **TF1 compatibility.** The new optimizer does not support TF1 any more,
        so please use the legacy optimizer `tf.keras.optimizer.legacy.XXX`.
        We highly recommend to migrate your workflow to TF2 for stable
        support and new features.