}
  }

  // Returns whether any tensor in a list of tensors is being watched and has
  // a trainable dtype.
  bool ShouldRecord(absl::Span<const int64_t> tensor_ids,
                    absl::Span<const tensorflow::DataType> dtypes) const;

  // Adds this tensor to the list of watched tensors.
  //
  // This is a no-op if the tensor is already being watched either from an

def TFR_CastOp : TFR_Op<"cast", [Pure]> {
  let description = [{
    The `cast` operation converts the operand with built-in tensor type to
    tfr.tensor type, or vice versa.

    Example:

    ```mlir
    %1 = tfr.cast(%0) : tensor<f32> -> !tfr.tensor
    %3 = tfr.cast(%1) : !tfr.tensor -> tensor<f32>
    ```
  }];

  let arguments = (ins TFR_singleTensorType:$arg);

                                       std::vector<Tensor>* x_datas,
                                       std::vector<Tensor>* jacobian_ts) {
  size_t y_num = y_shapes.size();
  size_t x_num = x_shapes.size();
  // x_stride and y_stride are used to calculate the correct jacobian row and
  // column position for a pair of elements at positions r, c within the x and y
  // tensors respectively.

) -> _SignatureDefMap:
  """Validates if the tensor names in signatures are consistent with the graph.

  This function checks if the input and output tensor names in the signatures
  exist if the graph. The output tensor names might change during conversion,
  we try to fix that with `_restore_output_tensor_names`. Besides, if there
  are duplicated tensor names, they we will be prefixed with the signature name.

class XlaComputationLaunchContext {
 public:
  // Create a new launch context. 'allocate_xla_tensors' is true if allocated
  // output tensors and variables are always XlaTensors. If false they are
  // assumed to be "normal" device pointers.
  // If 'use_multiple_streams' is true, tensors may be defined and used on
  // multiple streams and so se::Events must be defined and waited for. If

    const std::vector<int>& tensors = lstm_property.restrict_scale[0];
    if (tensors.size() != 2) {
      op.emitError(
          "Unexpected restricted_scale from operator property."
          " Should only have a pair of indices.");
      return failure();
    }
    return processRestrictScale(op, tensors[0], tensors[1], rewriter);
  }

 private:

    sess: Session instance used to evaluate tf.Tensors.

  Returns:
    The new representative dataset where each tf.Tensor is replaced by its
    evaluated numpy ndarrays.
  """
  new_repr_ds = []
  for sample in repr_ds:
    new_sample = {}
    for input_key, input_data in sample.items():
      # Evaluate the Tensor to get the actual value.
      if isinstance(input_data, core.Tensor):

        ) from ex


def _convert_values_to_tf_tensors(
    sample: rd.RepresentativeSample,
) -> Mapping[str, core.Tensor]:
  """Converts TensorLike values of `sample` to Tensors.

  Creates a copy of `sample`, where each value is converted to Tensors
  unless it is already a Tensor.
  The values are not converted in-place (i.e. `sample` is not mutated).

  Args:

  for (auto arg_num : input_mapping) {
    const Tensor* tensor;
    if (auto it = variable_snapshots.find(arg_num);
        it != variable_snapshots.end()) {
      tensor = it->second;
    } else {
      tensor = inputs[arg_num - num_missing_prefix_ctx_inputs];
    }

    // The input tensor can have the following cases.
    // 1. Tensor with PjRtTensorBuffer, containing a PjRtBuffer. This case

  // quantization aware training or calibration, for the remaining tensors.
  std::vector<std::pair<std::optional<double>, std::optional<double>>>
      input_ranges;

  // Whether to disable setting the quantization parameters of the input nodes
  // using input ranges.
  bool disable_set_input_nodes_quantization_params = false;

  // The default ranges can be used when a tensor doesn't have quantization