mlir::dyn_cast_or_null<ShapedType>(operand.getType());
          return result_type && operand_type && result_type.hasRank() &&
                 operand_type.hasRank() &&
                 result_type.getRank() == operand_type.getRank();
        });
    if (!is_valid_move) return;
  }

  // At this point we checked that we can safely move Transpose node before
  // `op`, and bypass all result transposes.

    for (const auto& axis_int : perm.getValues<APInt>()) {
      int64_t axis = axis_int.getSExtValue();
      if (axis < 0) {
        axis += input_type.getRank();
      }
      if (axis < 0 || (input_type.hasRank() && axis >= input_type.getRank())) {
        return transpose_op.emitOpError("perm must be in [-rank, rank)");
      }
      if (std::count(axes.begin(), axes.end(), axis) > 0) {

      BlockArgument arg = func.getArgument(index);
      auto inputType = mlir::dyn_cast<RankedTensorType>(arg.getType());
      // Only rank 1 tensor is supported for now.
      if (!inputType || inputType.getRank() != 1) continue;
      auto shape = llvm::to_vector<4>(inputType.getShape());
      llvm::SmallVector<int64_t, 4> bounds(shape.begin(), shape.end());
      // Mark the dim as dynamic dim.

    // Scalar identity is broadcastable to any operand shape, we only need to
    // check that operand has the same shape as a result.
    bool scalar_identity = identity_ty.hasRank() && identity_ty.getRank() == 0;
    if (scalar_identity) return operand_ty == result_ty;

    // If identity is not a scalar, we must verify that identity shape is
    // statically known to be broadcastable to the operand shape and the operand

      hlo::convertElementsAttr(group_assignment, builder.getIntegerType(64)));
  if (replica_groups.getType().getRank() != 2) {
    return op->emitOpError() << "group_assignment should have rank 2, got "
                             << replica_groups.getType().getRank();
  }
  return success();
}

ChannelHandleAttr ConvertChannel(OpBuilder& builder, int64_t channel_id,

// Checks if the value has rank at most 'n'.
class HasRankAtLeast<int n> : Constraint<
    CPred<"$0.getType().cast<ShapedType>().hasRank() && "
          "$0.getType().cast<ShapedType>().getRank() >= " # n>>;

// Checks value is not produced by a TFL_Quant or
// from TFL_Quant Op with same quant type.
def NotFromQuantOpOrSameQuantType : Constraint<
  CPred<"tflite::NotFromQuantOpOrSameQuantType($0,$1)">>;

    // Input should be rank 4.
    if (!input_ty.hasRank() || input_ty.getRank() != 4) {
      return failure();
    }

    // Check that out_size is rank-1, length-2. Otherwise the size is not legal.
    if (!out_size_ty.hasRank() || out_size_ty.getRank() != 1 ||
        out_size_ty.getShape()[0] != 2) {
      return failure();
    }

    // tensor, for setting depth_multiplier attribute, etc.).
    auto filter = tf_op.getFilter();
    auto filter_type = mlir::dyn_cast<RankedTensorType>(filter.getType());
    if (!filter_type || filter_type.getRank() != 4 ||
        !filter_type.hasStaticShape())
      return failure();

    Value input = tf_op.getInput();
    RankedTensorType input_type =
        mlir::dyn_cast<RankedTensorType>(input.getType());

       * layout using the TPU API. Running legalize_tf.cc on non-TPU nodes
       * thus is a potential source of bugs.
       */
      minor_to_major.resize(t.getRank());
      std::iota(minor_to_major.begin(), minor_to_major.end(), 0);
      std::sort(minor_to_major.begin(), minor_to_major.end(),
                [=](int64_t a, int64_t b) {
                  int64_t da = t.getDimSize(a);

//===----------------------------------------------------------------------===//

def GetBiasAddGradReductionIndices : NativeCodeCall<
  "GetBiasAddGradReductionIndices("
  "$0.getType().cast<RankedTensorType>().getRank(), $1, &$_builder)">;

def LowerBiasAddGradOp :
  Pat<(TF_BiasAddGradOp AnyRankedTensor:$out_backprop, $data_format),
      (TF_SumOp $out_backprop,
                (TF_ConstOp (GetBiasAddGradReductionIndices $out_backprop,