%0 = tfr.call @Huu(%arg0)  : (tensor<?xf32>) -> !tfr.tensor<K>
  func.return %0 : !tfr.tensor<K>
}

// -----

// CHECK-LABEL: get_shape
func.func @get_shape(%arg0: !tfr.tensor) -> (!shape.shape, !shape.shape) {
  %0 = tfr.get_shape %arg0 -> !shape.shape
  %1 = "tfr.get_shape"(%arg0) : (!tfr.tensor) -> !shape.shape
  func.return %0, %1 : !shape.shape, !shape.shape
}

// -----

// CHECK-LABEL: get_real_shape

  auto shape = mlir::cast<ShapedType>(value.getType());
  if (shape.getRank() != 1) {
    SmallVector<int64_t> new_shape;
    new_shape.push_back(shape.getNumElements());
    value = builder.create<TF::ReshapeOp>(
        loc, value, Create1DConstValue(builder, loc, new_shape));
  }
  return ConstantFoldOpIfPossible(value.getDefiningOp()).front();
}

  if (!type.hasRank()) return type;

  auto shape = type.getShape();
  SmallVector<int64_t, 4> new_shape(shape.size());

  for (int i = 0; i < permutation.size(); ++i) {
    int64_t index = permutation[i];
    assert(index < shape.size());
    new_shape[index] = shape[i];
  }

  return type.clone(new_shape);
}

// Move Transpose operations that permute `op` operands after the `op`.

      SmallVector<int64_t> new_shape(shape.begin(), shape.end());
      // If dimension of the input type is dynamic. Update the
      // bounds of the dim with the new type if needed.
      for (int i = 0; i < input_ty.getShape().size(); i++) {
        if (hlo::isDynamicDimSize(input_ty.getShape()[i])) {
          new_bounds[i] = new_shape[i];
          new_shape[i] = ShapedType::kDynamic;
        }
      }

                           Location loc) {
  auto buffer_type = mlir::cast<RankedTensorType>(buffer.getType());
  if (buffer_type.getShape().size() == 1) return index;
  // Create a concat of index and trailing zeros.
  llvm::SmallVector<int64_t, 8> zeros(buffer_type.getShape().size() - 1, 0);
  auto zeros_tensor = GetR1Const(zeros, builder, loc);
  return builder.create<TF::ConcatV2Op>(
      loc,

  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.

  // Extract Z from %dot.
  ArrayRef<int64_t> shape_z =
      dot.getType().getShape().drop_front(shape_b.size() + shape_y2.size());

  // Check %after shape.
  if (reshape_after.getType().getShape() !=
      ArrayRef<int64_t>(llvm::to_vector(
          llvm::concat<const int64_t>(shape_b, shape_y1, shape_z)))) {
    return failure();
  }

  rewriter.replaceOpWithNewOp<mhlo::DotGeneralOp>(

  RankedTensorType tensorType = mlir::cast<RankedTensorType>(value.getType());
  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

    }
  }];

  let hasCanonicalizer = 1;
}

def TFR_GetShapeOp : TFR_Op<"get_shape", [Pure]> {
  let description = [{
    The `get_shape` operation gets the shape of a tfr.tensor and returns
    !shape.shape type.

    Example:

    ```mlir
    %1 = "tfr.get_shape"(%0) : !tfr.tensor -> !shape.shape
    %1 = tfr.get_shape %0 -> !shape.shape
    ```
  }];

  let arguments = (ins TFR_TensorType:$arg);

                                   ConversionPatternRewriter& rewriter) {
  if (index_vector_dim == indices_type.getRank()) {
    llvm::SmallVector<int64_t, 4> new_start_indices_shape(
        indices_type.getShape().begin(), indices_type.getShape().end());
    new_start_indices_shape.push_back(1);
    indices_type = RankedTensorType::get(new_start_indices_shape,
                                         indices_type.getElementType());