if (auto tensor_type = dyn_cast<RankedTensorType>(input_type))
    return RankedTensorType::get(tensor_type.getShape(), elemental_type);
  if (auto tensor_type = dyn_cast<UnrankedTensorType>(input_type))
    return UnrankedTensorType::get(elemental_type);
  if (auto vector_type = dyn_cast<VectorType>(input_type))
    return VectorType::get(vector_type.getShape(), elemental_type);

  // If the expressed types match, just use the new elemental type.

  // CHECK: %[[LHS_2:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#1, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32>
  // CHECK: %[[LHS_3:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#2, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32>
  // CHECK: %[[LHS_4:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#3, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32>

    %0 = stablehlo.reshape %arg0 : (tensor<10x1x3xf32>) -> tensor<3x10xf32>
    return %0 : tensor<3x10xf32>
  }
  // CHECK: %[[RESHAPE:.*]] = stablehlo.reshape
  // CHECK-NEXT: return %[[RESHAPE]]

  // CHECK: @main_1
  func.func private @main_1(%arg0: tensor<3x10xf32>) -> tensor<6x5xf32> {
    %0 = stablehlo.reshape %arg0 : (tensor<3x10xf32>) -> tensor<6x5xf32>
    return %0 : tensor<6x5xf32>

// CHECK:           %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK:           %[[VAL_7:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>
// CHECK:           %[[VAL_8:.*]] = "tfl.reshape"(%[[VAL_2]], %[[VAL_4]]) : (tensor<384x384xf32>, tensor<4xi32>) -> tensor<1x1x384x384xf32>

  auto y = AvgPool3D(scope_, x, ksize, strides, "SAME");
  RunTest(x, x_shape, y, y_shape);
}

TEST_F(NNGradTest, LRN) {
  TensorShape x_shape({1, 1, 2, 1});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
  auto y = LRN(scope_, x);
  RunTest(x, x_shape, y, x_shape);
}

TEST_F(NNGradTest, SoftplusGrad) {
  TensorShape shape({3, 7});

ArrayRef<int64_t> GetRankedTensorShape(Value value) {
  return mlir::cast<RankedTensorType>(value.getType()).getShape();
}

Value SliceRankedTensor(OpBuilder* builder, Value input,
                        ArrayRef<int64_t> begin_shape,
                        ArrayRef<int64_t> begin_values,
                        ArrayRef<int64_t> size_shape,
                        ArrayRef<int64_t> size_values,
                        mlir::Location location) {

                             input_vector_type.getElementType();
    bool same_shape =
        output_tensor_type.getShape() == input_vector_type.getShape();
    if (!same_element_type || !same_shape) {
      op.emitError("input and output should have same shape and element type.");
    }
    return success(same_element_type && same_shape);
  }

  op.emitError("input can not be converted to an output tensor.");

  if (!dim_nums.getLhsBatchingDimensions().empty()) return failure();

  auto reshape = dot.getRhs().getDefiningOp<mhlo::ReshapeOp>();
  if (!reshape) return failure();
  if (!reshape->hasOneUse())
    return rewriter.notifyMatchFailure(reshape, "reshape has multiple usages");
  if (!reshape.getType().hasStaticShape() ||
      !reshape.getOperand().getType().hasStaticShape() ||
      !dot.getLhs().getType().hasStaticShape()) {

#define TENSORFLOW_C_TF_SHAPE_INTERNAL_H_

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

typedef struct TF_Shape TF_Shape;

namespace tensorflow {

DEFINE_CONVERSION_FUNCTIONS(tensorflow::PartialTensorShape, TF_Shape);

}

  Type element_type = tensorType.getElementType();

  int rank = tensorType.getShape().size();
  int num_rows = tensorType.getShape()[rank - 2];
  int num_cols = tensorType.getShape()[rank - 1];

  std::vector<Value> sliced;

  if (batch_size == 1) {
    // Batch size is 1, no splitting is required
    // Squeeze the batch dimension, i.e. reshape
    // [1, num_rows, num_cols] -> [num_rows, num_cols]