}
        auto attr_value =
            mlir::DenseI32ArrayAttr::get(inner_op->getContext(), values);
        inner_op->setAttr(attr_name, attr_value);
      };

  if (inner_op->hasTrait<mlir::OpTrait::AttrSizedOperandSegments>() ||
      inner_op->hasTrait<mlir::OpTrait::AttrSizedResultSegments>()) {
    // The op has multiple variadic operands or results.

// Tests that a quantized `stablehlo.pad` without interior padding is
// converted to `tfl.padv2`.

func.func @pad_without_interior_padding(
    %arg0: tensor<2x3x!quant.uniform<i8:f32, 2.000000e+00:-1>>,
    %arg1: tensor<!quant.uniform<i8:f32, 2.000000e+00:-1>>
  ) -> tensor<4x5x!quant.uniform<i8:f32, 2.000000e+00:-1>> {
  %0 = stablehlo.pad %arg0, %arg1, low = [0, 1], high = [2, 1], interior = [0, 0] : (

        /*fused_activation_function=*/rewriter.getStringAttr("NONE"));
  }
};

// Rewrites quantized stablehlo.pad to tfl.padv2.
// tfl.dilate is introduced in between when interior padding exists.
// TODO: b/322428814 - Add StableHLO quantizer integration tests for ODML.
class RewriteQuantizedPadOp : public OpRewritePattern<stablehlo::PadOp> {
 public:

      control_nodes_at.try_emplace(&item.value(), item.index());
    }
  }

  for (auto outer_op : control_nodes) {
    auto control_node_op = dyn_cast<mlir::TFL::ControlNodeOp>(outer_op);
    auto* inner_op = &control_node_op.getBody().front().front();
    auto control_token = control_node_op.getControl();

    // Now go through all uses. Since *block is in executable order, control

// Returns if the shape inference pass supports an op outside the TF dialect.
bool IsSupportedNonTFOp(Operation* op) {
  return isa<tf_device::ReturnOp, tf_device::ClusterOp, tf_device::LaunchOp,
             tf_executor::EnterOp, tf_executor::ExitOp, tf_executor::FetchOp,
             tf_executor::GraphOp, tf_executor::IslandOp,
             tf_executor::LoopCondOp, tf_executor::MergeOp,