// 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: %[[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 ||

  strides = [1, op.get_attr('stride_w'), op.get_attr('stride_h'), 1]
  padding = op.get_attr('padding')
  shape_0, shape_1 = tf.shape_n([op.inputs[0], op.inputs[1]])
  return [
      tf.compat.v1.nn.conv2d_backprop_input(
          shape_0,
          op.inputs[1],
          grad,
          strides=strides,
          padding=padding,
          dilations=dilations,

     (HasNoUseOf:$root__3), (HasNoUseOf:$root__4),
     (HasNoUseOf:$root__5)]>;

class HasEqualElementSize<list<int> shape_1, list<int> shape_2> : Constraint<
  CPred<"quant::HasEqualElementSize($0, $1,"
  "llvm::ArrayRef<int>({" # !interleave(shape_1, ", ") # "}),"
  "llvm::ArrayRef<int>({" # !interleave(shape_2, ", ") # "}))">,
  "Checks if the given dimensions contain the same number of elements.">;

def ReshapableTo1DTensor : Constraint<

    OpBuilder builder(shape_op);
    auto new_shape_op = builder.create<TF::VariableShapeOp>(
        shape_op.getLoc(), shape_op.getType(),
        replicate_op.GetReplicaOperandForBlockArgument(block_arg,
                                                       /*replica=*/0));
    shape_op.replaceAllUsesWith(new_shape_op.getOperation());
    shape_op.erase();
  }
}

          TF::IdentityOp,
          TF::MaxPoolOp,
          TF::PadV2Op,
          TF::RankOp,
          TF::ReshapeOp,
          TF::SelectOp,
          TF::SelectV2Op,
          TF::ShapeNOp,
          TF::ShapeOp,
          TF::SizeOp,
          TF::SqueezeOp,
          TF::TransposeOp
          // go/keep-sorted end
          // clang-format on
          >(op)) {

TensorSpec::TensorSpec()
    : shape_(std::initializer_list<int64_t>()), dtype_(DT_FLOAT) {}

TensorSpec::TensorSpec(PartialTensorShape shape, DataType dtype)
    : shape_(std::move(shape)), dtype_(dtype) {}

TensorSpec::TensorSpec(const TensorSpecProto& proto)
    : shape_(proto.shape()), dtype_(proto.dtype()) {}

const PartialTensorShape& TensorSpec::shape() const { return shape_; }

      return;
    }
  }

  // These ops are cheap regardless of their input sizes.
  //
  // TODO(chky): Find a more scalable way to figure out cheap ops.
  if (llvm::isa<mlir::TF::ShapeOp, mlir::TF::StridedSliceOp,
                mlir::TF::ReshapeOp, mlir::TF::ExpandDimsOp>(op)) {
    cost_map_[op] = kDefaultCheapCost;
    return;
  }

  // For other ops, use the sum of input sizes as its cost.

TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b);

// Return an identity op.
TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return a shape op fetching the shape of `a`.
TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return an allreduce op adding up input tensor `in` from `group_size` workers.
TFE_Op* AllReduceOp(TFE_Context* ctx, TFE_TensorHandle* in, int group_size);