// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[CONV]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK: %[[SHAPE_OF_1:.*]] = shape.shape_of %[[ADD]]
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>

  // CHECK-DAG:  %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:  %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1
  // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]]
  // CHECK-NEXT: shape.assuming %[[WITNESS:.+]]
  // CHECK-DAG:    %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:    %[[RHS_SHAPE:.+]] = shape.shape_of %arg1

// CHECK-LABEL:   func @batchmatmulv2_basic
// CHECK-SAME:        ([[LHS:%.*]]: tensor<1x4x2xf32>, [[RHS:%.*]]: tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
// CHECK:           [[LHSSHAPE:%.*]] = shape.shape_of [[LHS]] : tensor<1x4x2xf32>
// CHECK:           [[RHSSHAPE:%.*]] = shape.shape_of [[RHS]] : tensor<3x2x4xf32>
// CHECK:           [[CM2:%.*]] = arith.constant -2 : index
// CHECK:           [[LHSHEAD:%.*]], [[LHSTAIL:%.*]] = "shape.split_at"([[LHSSHAPE]], [[CM2]])

// CHECK-DAG: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] : tensor<?x3x2x2xf32> -> tensor<4xindex>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST_0]], %[[SHAPE_OF]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x3x2x2xf32>

  // CHECK: %[[SHAPE:.+]] = shape.shape_of %[[CONV]] : tensor<?x256x256x2xf32> -> tensor<4xindex>

        diag << "entities 'conv' failed to satisfy constraint: has one use";
      });
    }
    // For dynamic case, the result of conv should be used by shape_of and mul.
    if (is_dynamic_broadcast) {
      auto conv_uses = (*conv_op.getODSResults(0).begin()).getUses();
      if (std::distance(conv_uses.begin(), conv_uses.end()) != 2 ||

    %mean: tensor<?xf32>, %variance: tensor<?xf32>)
    -> tensor<?x?x?x?xf32> {
  // CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e-03> : tensor<f32>
  // CHECK-DAG: %[[VAR_SHAPE:.+]] = shape.shape_of %[[VARIANCE]] : tensor<?xf32> -> tensor<1xindex>

    %2 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<?x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>
    %3 = shape.shape_of %2 : tensor<?x2xf32> -> tensor<2xindex>
    %4 = stablehlo.dynamic_broadcast_in_dim %arg2, %3, dims = [1] : (tensor<2xf32>, tensor<2xindex>) -> tensor<?x2xf32>
    %5 = stablehlo.add %2, %4 : tensor<?x2xf32>

  // CHECK-DAG: %[[SHAPEOF:.+]] = shape.shape_of %arg0
  // CHECK-DAG: %[[CST0:.+]] = arith.constant 0
  // CHECK-DAG: %[[CST1:.+]] = arith.constant 1
  // CHECK-DAG: %[[GETEXTENT0:.+]] = tensor.extract %[[SHAPEOF]][%[[CST0]]]
  // CHECK-DAG: %[[CST1_0:.+]] = arith.constant 1
  // CHECK-DAG: %[[GETEXTENT1:.+]] = tensor.extract %[[SHAPEOF]][%[[CST1_0]]]

 public:
  LogicalResult matchAndRewrite(GetShapeOp shape_op,
                                PatternRewriter &rewriter) const override {
    Operation *preceding_op = shape_op.getArg().getDefiningOp();
    if (auto cast_op = llvm::dyn_cast_or_null<CastOp>(preceding_op)) {
      // replace this pair by shape.shape_of, so the folding works.
      rewriter.replaceOpWithNewOp<shape::ShapeOfOp>(shape_op, cast_op.getArg());
      return success();
    }