_, reduction_axes = tf.raw_ops.BroadcastGradientArgs(
      s0=broadcast_shape, s1=input_value_shape)
  updates_grad_reshaped = tf.reduce_sum(
      grad, axis=reduction_axes, keepdims=True)
  bias_grad = tf.reshape(updates_grad_reshaped, input_value_shape)

  dilations = [1, op.get_attr('dilation_w'), op.get_attr('dilation_h'), 1]
  strides = [1, op.get_attr('stride_w'), op.get_attr('stride_h'), 1]
  padding = op.get_attr('padding')

// CHECK-NEXT: [[o8:%.*]] = tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.Reshape"([[o7]], [[o1]])
// CHECK-NEXT: [[o9:%.*]] = tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf._FusedMatMul"([[o8]], [[o5]], [[o4]])

  auto input_type = mlir::cast<mlir::TensorType>(src_input.getType());

  if (input_type.hasRank()) {
    if (input_type.getShape()[split_dimension] == mlir::ShapedType::kDynamic) {
      output_type = input_type;
    } else {
      auto shape = llvm::to_vector<4>(input_type.getShape());
      if (shape[split_dimension] % num_split != 0) {
        return mlir::emitError(
            location,
            llvm::formatv(

    %22 = "tfl.concatenation"(%cst_6, %21) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1xi32>, tensor<?xi32>) -> tensor<?xi32>
    %23 = "tfl.reshape"(%arg2, %cst_12) : (tensor<i32>, tensor<1xi32>) -> tensor<1xi32>
    %24 = "tfl.fill"(%cst_11, %cst_9) : (tensor<1xi32>, tensor<i32>) -> tensor<?xi32>

//
// The current heuristic used is quite simple, which is to calculate the total
// size of input tensors. The exception is that ops whose cost is irrelevant to
// input sizes, such as tf.Shape and tf.Reshape, are whitelisted to have cheap
// cost. This cost analysis is expected to be used conservatively (eg. use a low
// threshold to decide whether a cost is cheap or expensive), as it might not be
// accurate in some cases.
//

    (Arith_ConstantOp:$permutation2 $p2)),
  (TF_TransposeOp $input,
    (Arith_ConstantOp (RemapPermutation $permutation1, $permutation2))),
  [(HasOneUse $transpose_out1)]>;

// Pattern to fuse trivial reshape op into transpose op
def FoldTrivialReshapeIntoTranspose : Pat<
  (TF_ReshapeOp:$output
      (TF_TransposeOp:$transpose_out1 $input, (Arith_ConstantOp:$permutation1 $p1)), $_),
  (TF_TransposeOp:$transpose_op $input,

tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>) {\0A    %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) -> tensor<3x1120x?xi32> loc(#loc9)\0A    %1:3 = \22tf.Split\22(%arg2, %0) {_xla_outside_compilation = \220\22} : (tensor<i32>, tensor<3x1120x?xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1x1120x?xi32>) loc(#loc10)\0A    %2 = \22tf.Reshape\22(%1#0, %arg3) {_xla_outside_compilation = \220\22} : (tensor<1x1120x?xi32>,...

  %9 = mhlo.compare  LT, %8, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
  %10 = mhlo.add %8, %4 : tensor<32x32xi64>
  %11 = mhlo.select %9, %10, %8 : tensor<32x32xi1>, tensor<32x32xi64>
  %12 = mhlo.reshape %11 : (tensor<32x32xi64>) -> tensor<32x32x1xi64>
  %13 = "mhlo.broadcast_in_dim"(%7) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<32xi64>) -> tensor<32x32xi64>

        "@local_xla//xla/tsl/c:tsl_status_internal_headers",
    ],
    visibility = [
        "//tensorflow/python:__subpackages__",
    ],
)

cc_library(
    name = "tf_shape",
    srcs = ["tf_shape.cc"],
    hdrs = ["tf_shape.h"],
    copts = tf_copts(),
    visibility = ["//visibility:public"],
    deps = [
        ":c_api_macros",
        ":tf_shape_internal",
        "//tensorflow/core:framework",

  %8 = mhlo.add %6, %7 : tensor<32x28x28x5xf32>
  %9 = mhlo.maximum %8, %1 : tensor<32x28x28x5xf32>
  %10 = "mhlo.reshape"(%9) : (tensor<32x28x28x5xf32>) -> tensor<32x3920xf32>
  %11 = "mhlo.dot"(%10, %5) : (tensor<32x3920xf32>, tensor<3920x10xf32>) -> tensor<32x10xf32>