%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

    }
  }];

  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);

// Tests for ops with canonicalization patterns.

// CHECK-LABEL: get_real_shape
func.func @get_real_shape(%arg0: tensor<1x2xf32>) -> tensor<2xindex> {
  %0 = "tfr.cast"(%arg0) : (tensor<1x2xf32>) -> !tfr.tensor
  %1 = tfr.get_shape %0 -> !shape.shape
  %2 = shape.to_extent_tensor %1 : !shape.shape -> tensor<2xindex>
  func.return %2 : tensor<2xindex>

// CHECK-NEXT: %[[e:.*]] = shape.const_shape [1, 2] : tensor<2xindex>

    mlir_code_exp = r"""
      CHECK-LABEL: tfr.func @tf__test_identity_op(%x: !tfr.tensor) -> (!tfr.tensor) {
      CHECK-NEXT:   %[[shape:.*]] = tfr.get_shape %x -> !shape.shape

      CHECK-NEXT:   %[[shape_1:.*]] = tfr.get_shape %x -> !shape.shape
      CHECK-NEXT:   %[[len:.*]] = shape.rank %[[shape_1]] : !shape.shape -> !shape.size
      CHECK-NEXT:   %[[index:.*]] = shape.size_to_index %[[len]] : !shape.size

  }
  int rank = shaped_type.getRank();
  if (rank < 2) {
    return nullptr;
  }
  SmallVector<int64_t> new_shape(shaped_type.getShape().begin(),
                                 shaped_type.getShape().end());
  std::swap(new_shape[rank - 1], new_shape[rank - 2]);
  return shaped_type.clone(new_shape);
}

// Returns a ShapedType for a permutation and the shape of input after

              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];
    }
  }

    // Get the new shape.
    SmallVector<int64_t, 4> new_shape;
    for (int i = 0; i < 4 - rank; ++i) {
      new_shape.push_back(1);
    }
    for (auto size : input_type.getShape()) {
      new_shape.push_back(size);
    }

    auto reshape_op =
        InsertReshapeOp(concat_op.getLoc(), input, input_type.getElementType(),
                        new_shape, &rewriter);
    reshape_ops.push_back(reshape_op.getResult());

            mlir::cast<ShapedType>(lhs.getType()).getShape(),
            mlir::cast<ShapedType>(rhs.getType()).getShape(),
            symbolic_broadcast_shape)) {
      return failure();
    }

    // Calculates the broadcast shape using BroadcastArgs op.
    Value lhs_shape = GetShape(lhs, op->getLoc(), rewriter);
    Value rhs_shape = GetShape(rhs, op->getLoc(), rewriter);
    auto broadcast_shape =

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));

      const int64_t b = lhs_type.getShape()[std::get<0>(i)];
      lhs_shape.push_back(b);
      out_shape->push_back(b);

      lhs_reshape_segids.at(std::get<0>(i)) = num_lhs_reshape_segids++;
    }
  } else {
    const int64_t lhs_out_size =
        ProdShapeWithIndexInTuple<0>(lhs_type.getShape(), dnums.lhs_out);
    lhs_shape.push_back(lhs_out_size);
    out_shape->push_back(lhs_out_size);