}
  }

  ArrayRef<int64_t> shape = ranked_ty.getShape();
  SmallVector<int64_t, 4> out_shape;
  out_shape.reserve(rank - (keep_dims.getValue() ? 0 : num_reduce_dim));
  for (int64_t i = 0; i < rank; ++i) {
    if (!is_reduce_dim[i])
      out_shape.push_back(shape[i]);
    else if (keep_dims.getValue())
      out_shape.push_back(1);
  }
  return RankedTensorType::get(out_shape, element_ty);
}

    ArrayRef<int64_t> shape = ranked_type.getShape();
    assert(permutation.size() == shape.size());

    SmallVector<int64_t, 4> new_shape(permutation.size());
    for (size_t i = 0; i < permutation.size(); ++i)
      new_shape[i] = shape[permutation[i]];

    return RankedTensorType::get(new_shape, ranked_type.getElementType());
  }

  return type;
}

    %2 = "tf.PartitionedCall"(%1) {config = "", config_proto = "", executor_type = "", f = @get_shape} : (tensor<?x?xf32>) -> (tensor<?xi64>)

    // CHECK: return %[[RESULT]]
    func.return %2 : tensor<?xi64>
  }

  // CHECK-LABEL: func @get_shape
  func.func @get_shape(%arg0 : tensor<*xi64>) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<*xi64>) -> tensor<?xi64>

  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.

def GetSqueezedPermutation: NativeCodeCall<"GetSqueezedPermutation($0, $1)">;

// Check to see if the tensor dimensions can be Squeezed by eliminating 1s'
def CanSqueezeTensor : Constraint<CPred<
  "GetShape($0).getNumElements() > GetSqueezedShape($0).getNumElements()">>;


// Pattern to convert TFL_TransposeOp with rank>6 to rank<=6 if there are
// redundant dimensions in the tensor. For example- [2x1x3] == [2x3] and 1 is

      "outputs": [1],
      "operators": [
        {
          "inputs": [ 0 ],
          "outputs": [ 1 ],
          "builtin_options_type": "ReshapeOptions",
          "builtin_options": {
            "new_shape": [ 2, 2 ]
          }
        }
      ]
    }
  ],
  "buffers": []

    auto rhs_type = mlir::dyn_cast<RankedTensorType>(input_rhs.getType());

    if (!lhs_type || !rhs_type) return failure();

    auto lhs_shape = lhs_type.getShape();
    auto rhs_shape = rhs_type.getShape();

    // Ensure that input ranks are at least 2.
    const int dims_a = lhs_shape.size();
    const int dims_b = rhs_shape.size();
    if (dims_a < 2 || dims_b < 2) {
      return failure();

    // Get the unbroadcasted shapes in the operand order.
    std::array<llvm::ArrayRef<int64_t>, 2> operand_shapes;
    operand_shapes[i] = broadcast_arg_type.getShape();
    operand_shapes[1 - i] = argument_type.getShape();

    // Check that the input of the broadcast and the other operand is broadcast
    // compatible.
    llvm::SmallVector<int64_t, 4> broadcasted_shape;

def CreateNoneValue : NativeCodeCall<
  "$_builder.create<TFL::NoValueOp>($0.getLoc(), $_builder.getUnitAttr())">;

// Returns shape of a ranked tensor.
// if called without a ranked tensor it will fail.
def GetShape: NativeCodeCall<"GetShape($0)">;

// Constraint that values in list attribute are all ones.
def IsAllOnesConstant : Constraint<CPred<"TFL::IsAllOnesConstant($0)">>;

  ShapedType packed_shape_type =
      mlir::dyn_cast<ShapedType>(packed_value.getType());
  llvm::SmallVector<int64_t> packed_shape(packed_shape_type.getShape().begin(),
                                          packed_shape_type.getShape().end());
  EXPECT_THAT(packed_shape, testing::ElementsAreArray(expected_packed_shape));
  llvm::SmallVector<int8_t> packed_value_vector(
      packed_value_attr.getValues<int8_t>());