TF_RETURN_IF_ERROR(
      tensorflow::TensorShapeToXLAShape(type, shape, &xla_shape));
  ApiConverter::StackHelper<XLA_Shape> se_shape(xla_shape);
  ApiConverter::StackHelper<XLA_Shape> tpu_shape;
  StatusHelper status;
  stream_executor::tpu::ExecutorApiFn()->XlaShapeToTpuShapeRepresentationFn(
      &se_shape.value, type, use_fast_memory, &tpu_shape.value,
      status.c_status);
  if (!status.status().ok()) {
    return status.status();

  // If the reshape changes the number of dimensions so it cannot be interpreted
  // as a transpose.
  if (in_shape.size() != out_shape.size()) {
    return false;
  }

  in_shape.erase(std::remove(in_shape.begin(), in_shape.end(), 1),
                 in_shape.end());
  out_shape.erase(std::remove(out_shape.begin(), out_shape.end(), 1),
                  out_shape.end());
  return in_shape == out_shape;
}

#endif

// An opaque type corresponding to a shape in tensorflow. In the future,
// we may expose the ABI of TF_Shape for performance reasons.
typedef struct TF_Shape TF_Shape;

// Return a new, unknown rank shape object. The caller is responsible for
// calling TF_DeleteShape to deallocate and destroy the returned shape.
TF_CAPI_EXPORT extern TF_Shape* TF_NewShape();

// Returns the rank of `shape`. If `shape` has unknown rank, returns -1.

#include "tensorflow/c/tf_shape.h"

#include <stdint.h>

#include "tensorflow/c/tf_shape_internal.h"
#include "tensorflow/core/framework/tensor_shape.h"

extern "C" {

TF_Shape* TF_NewShape() {
  return tensorflow::wrap(new tensorflow::PartialTensorShape());
}

int TF_ShapeDims(const TF_Shape* shape) {
  return tensorflow::unwrap(shape)->dims();
}

  if (s0_shape == s1_shape) return;

  for (int i = bcasted_shape.size(); i > 0; --i) {
    int idx = bcasted_shape.size() - i;
    int s0_idx = i > s0_shape.size() ? -1 : s0_shape.size() - i;
    int s1_idx = i > s1_shape.size() ? -1 : s1_shape.size() - i;
    if (s0_idx == -1) {
      r0.push_back(idx);
      if (s1_shape[s1_idx] == 1) r1.push_back(idx);
    } else if (s1_idx == -1) {

    // only when height = width.
    if (params_type.getShape().size() != 4 ||
        indices_type.getShape().size() != 2)
      return failure();
    if (params_type.getShape()[1] != 1) return failure();
    if (params_type.getShape()[0] != params_type.getShape()[2])
      return failure();
    if (result_type.getShape()[0] != params_type.getShape()[0] * 2)
      return failure();

  if (!new_size) return failure();

  // Reshape the input to 4d.
  SmallVector<int64_t, 4> new_shape;
  const int current_rank = input_type.getRank();
  for (int i = 0; i < 4 - current_rank; ++i) {
    new_shape.push_back(1);
  }
  for (auto size : input_type.getShape()) {
    new_shape.push_back(size);
  }

  auto reshape_op =

}

// Creates reshape op for space to depth transform.
TF::ReshapeOp GetReshapeOpForConv2DFilter(ArrayRef<int64_t> new_shape,
                                          Value input, OpBuilder* builder) {
  auto reshape_result_type =
      RankedTensorType::get(new_shape, getElementTypeOrSelf(input));
  auto reshape_type = RankedTensorType::get(
      {static_cast<int64_t>(new_shape.size())}, builder->getIntegerType(64));

  // https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/batch-mat-mul
  int64_t x_row_dim = x_shape[x_shape.size() - 2];
  int64_t x_col_dim = x_shape[x_shape.size() - 1];
  int64_t y_row_dim = y_shape[y_shape.size() - 2];
  int64_t y_col_dim = y_shape[y_shape.size() - 1];
  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

              mlir::cast<ShapedType>(slice_sizes.getType()).getShape(),
              builder.getI64Type()),
          slice_sizes));

  // Collapses dimensions by reshaping.
  SmallVector<int64_t> new_shape(operand_rank - collapsed_dims.size());
  for (int64_t i = 0, j = 0; i < operand_rank; i++) {
    if (!collapsed_dims.contains(i)) {
      new_shape[j++] = operand_shape[i];
    }
  }