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.

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

  TorchAvgPoolData data;

  auto op_type = mlir::cast<RankedTensorType>(op.getOperand(0).getType());

  data.n = op_type.getShape()[0];
  data.c = op_type.getShape()[1];
  data.h_in = op_type.getShape()[2];
  data.w_in = op_type.getShape()[3];

  std::vector<int32_t> kernel_size;
  GetI32VectorFromDenseI64CompositeAttr(composite_attrs, "kernel_size",

    // Thus, we fail to match if the consuming reshape rank is larger.
    ArrayRef<int64_t> input_shape = input_type.getShape();
    if (reshape_shape.size() > input_shape.size()) return failure();

    // Extend the input shape with leading 1s to match the broadcast shape.
    ArrayRef<int64_t> broadcast_shape = output_type.getShape();
    SmallVector<int64_t, 4> input_shape_extended;

          bmm_op->getLoc(), permuation_tensor_type,
          DenseElementsAttr::get(permuation_tensor_type, permute));

      auto input_shape = input_type.getShape();
      llvm::SmallVector<int64_t, 4> permuted_shape(input_shape.begin(),
                                                   input_shape.end());
      // Swaps z dimension and x dimension to get permuted shape.

// ShapeContainerT is any type with the following methods:
//   bool hasRank()
//   ArrayRef<int64_t> getShape()
// This includes mlir::TF::ShapeAttr and mlir::ShapedType.
template <typename ShapeContainerT>
void SetTensorShapeProto(ShapeContainerT shape, TensorShapeProto* proto) {
  if (shape.hasRank()) {
    for (int64_t dim : shape.getShape()) {
      proto->add_dim()->set_size(mlir::ShapedType::isDynamic(dim) ? -1 : dim);
    }
  } else {

    return operand_ty == result_ty && identity_ty.hasStaticShape() &&
           result_ty.hasStaticShape() &&
           OpTrait::util::staticallyKnownBroadcastable(operand_ty.getShape(),
                                                       identity_ty.getShape());
  };

  // Check that we have a constant operand on one side (candidate for identity).
  const bool is_commutative =