}

// CHECK-LABEL: func @testPow
// CHECK-SAME:(%[[ARG_0:.*]]: tensor<4xf32>, %[[ARG_1:.*]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>)
func.func @testPow(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) {

  %cst_zero = arith.constant dense<0.0> : tensor<f32>
  %cst_one = arith.constant dense<1.0> : tensor<f32>

}

# MLIR-LABEL: func @main
# MLIR-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<1x1x1x256x!quant.uniform<i8:f32, 0.21632751372549019:27>>) -> tensor<1x6x31x!quant.uniform<i8:f32, 0.09363494573854933:22>>
# MLIR-SAME:  control_outputs = ""
# MLIR-SAME:  inputs = "input"
# MLIR-SAME:  outputs = "output"
# MLIR:         %[[shape:.*]] = arith.constant dense<[1, -1, 31]> : tensor<3xi32>

    "mhlo.return"(%0, %1) : (tensor<f32>, tensor<f32>) -> ()
  },  {
    "mhlo.return"(%0, %1) : (tensor<f32>, tensor<f32>) -> ()
  // CHECK:      (tensor<i1>) -> (tensor<f32>, tensor<f32>, !mhlo.token)
  }) : (tensor<i1>) -> (tensor<f32>, tensor<f32>)

  // CHECK:      [[IF_SUBTUPLE_RESULT:%.*]] = mhlo.tuple [[IF]]#0, [[IF]]#1

      return tflite::CreateBuffer(builder_, buffer_data);
    }
  }

  tensorflow::Tensor tensor;
  auto status = tensorflow::ConvertToTensor(attr, &tensor);
  if (!status.ok()) {
    inst->emitError(
        Twine("failed to convert value attribute to tensor with error: " +
              status.ToString()));
    return std::nullopt;
  }

func.func @testBroadcastGradientArgs(%s0: tensor<4xi32>, %s1: tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>) {
  %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) : (tensor<4xi32>, tensor<4xi32>) -> (tensor<1xi32>, tensor<0xi32>)
  func.return %r0, %r1 : tensor<1xi32>, tensor<0xi32>
}

// -----

func.func @testBroadcastGradientArgsIncompatibleInputType(%s0: tensor<4xi32>, %s1: tensor<4xi64>) -> (tensor<1xi32>, tensor<0xi32>) {

  // copy-on-read mode is false.
  if (var->tensor()->RefCountIsOne()) {
    var->copy_on_read_mode.store(true);
    return absl::OkStatus();
  }
  Tensor tmp;
  if (variantType) {
    AllocatorAttributes attr;
    attr.set_on_host(true);
    TF_RETURN_IF_ERROR(context->allocate_temp(
        var->tensor()->dtype(), var->tensor()->shape(), &tmp, attr));

    const auto elements_in = var->tensor()->flat<Variant>();

    // Input tensors
    const std::vector<std::pair<string, Tensor>>& input_pairs,
    // Output tensors
    const std::vector<string>& output_tensor_names, TF_Tensor** c_outputs,
    // Target nodes
    const std::vector<string>& target_oper_names, TF_Buffer* run_metadata,
    TF_Status* status) {
  const int noutputs = output_tensor_names.size();
  std::vector<Tensor> outputs(noutputs);
  Status result;

// exporting passes is controlled by `export_opts`. Returns `AssetFileDef`s that
// associate the input arguments of @main and the asset file names. Asset file
// names will be used to feed the corresponding tensors during initialization
// upon model loading.
// TODO: b/329206105 - Add unit tests after decomposing post processing passes.
absl::StatusOr<SmallVector<::tensorflow::AssetFileDef>> RunExportPasses(

      ctx, inputs[0], inputs[1], &loss, &backprop,
      "SparseSoftmaxCrossEntropyWithLogits"));
  // `gradient_checker` only works with model that returns only 1 tensor.
  // Although, `ops::SparseSoftmaxCrossEntropyWithLogits` returns 2 tensors, the
  // second tensor isn't needed for computing gradient so we could safely drop
  // it.
  outputs[0] = loss;
  backprop->Unref();
  return absl::OkStatus();
}

// on-device shape. This is accurate for CPU and GPU devices that neither
// transpose nor pad tensors.
Status DefaultPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) {
  const tensorflow::XlaTensor* xla_tensor =
      tensorflow::XlaTensor::FromTensor(&tensor);
  if (xla_tensor == nullptr) {
    return TensorShapeToXLAShape(tensor.dtype(), tensor.shape(), shape);
  }

  const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer();