// 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>

    mlir::DenseIntElementsAttr shape_attr;
    if (matchPattern(op_state.operands[1], m_Constant(&shape_attr))) {
      auto shape_ty =
          op_state.operands[1].getType().dyn_cast<RankedTensorType>();
      if (shape_ty != nullptr && shape_ty.hasRank() && shape_ty.getRank() > 1) {
        llvm::SmallVector<mlir::Attribute, 4> shape;
        int32_t dim_size = 0;
        for (const auto& dim :

    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());

    TFE_Op* shape_op = ShapeOp(ctx, hgpu);
    TFE_OpSetDevice(shape_op, gpu_device_name.c_str(), status.get());
    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
    TFE_TensorHandle* retvals[1];
    int num_retvals = 1;
    TFE_Execute(shape_op, &retvals[0], &num_retvals, status.get());

          loc, GetScalarOfType(
                   builder->getIntegerType(use_32bit.getValue() ? 32 : 64),
                   ranked_ty.getDimSize(idx)));
    }
  }

  auto shape = builder->create<TF::ShapeOp>(loc, input, use_32bit);
  return builder->create<TF::StridedSliceOp>(
      loc, mlir::RankedTensorType::get({}, getElementTypeOrSelf(shape)), shape,
      /*begin=*/

  std::string serialized_func_module;
  if (HasDynamicOutputs(external_outputs)) {
    FuncOp shape_op = BuildFunction(clustered_ops, external_operands,
                                    external_outputs, &builder);
    EncapsulateFuncAndSerialize(shape_op, &serialized_func_module);
  }

  builder.setInsertionPoint(&op);
  SmallVector<Value, 4> host_compute_outputs;