llvm::SmallVector<int64_t>{});
    mhlo::BroadcastInDimOp broadcast_in_dim_op =
        op_builder->create<mhlo::BroadcastInDimOp>(
            const_op->getLoc(), splat_elements_attr.getType(), scalar,
            broadcast_dims);
    const_op->replaceAllUsesWith(broadcast_in_dim_op);
    const_op->erase();
  }
};

}  // namespace

    Value subtract_op_second_operand = subtract_op->getOperand(1);
    if (auto broadcast_in_dim_op = TryCast<stablehlo::BroadcastInDimOp>(
            subtract_op_second_operand.getDefiningOp(),
            /*name=*/"broadcast_in_dim_for_subtract_op_operand");
        succeeded(broadcast_in_dim_op)) {
      subtract_op_second_operand = broadcast_in_dim_op->getOperand();
    }

// `BroadcastInDimOp`'s operand is a block argument.
absl::Nullable<Operation*> SkipUpwardsOptionalBroadcastInDimOp(
    absl::Nonnull<Operation*> op) {
  if (auto broadcast_in_dim_op = dyn_cast_or_null<BroadcastInDimOp>(op);
      broadcast_in_dim_op != nullptr) {
    return broadcast_in_dim_op.getOperand().getDefiningOp();
  }
  return op;
}

class DeferActivationTransposeForAddOp : public OpRewritePattern<AddOp> {
 public:

// ```
template <typename T>
Operation* GetBroadcastedUserOp(Operation* op) {
  // Broadcast bias for known input shape.
  auto broadcast_in_dim_op = FindUserOfType<BroadcastInDimOp>(op);
  if (broadcast_in_dim_op != nullptr) {
    auto target_op = FindUserOfType<T>(broadcast_in_dim_op);
    if (target_op != nullptr) return target_op;
  }
  // Broadcast bias for unknown input shape.

// If no such constant op exists, returns a nullptr.
Operation* GetBiasConstOp(Operation* op) {
  Operation* bias_const_op;
  if (Operation* broadcast_in_dim_op =
          FindOperandOfType<stablehlo::BroadcastInDimOp>(op);
      broadcast_in_dim_op != nullptr) {
    bias_const_op = broadcast_in_dim_op->getOperand(0).getDefiningOp();
  } else {
    bias_const_op = FindOperandOfType<stablehlo::ConstantOp>(op);
  }

  %0 = "vhlo.broadcast_in_dim_v1"(%arg0) <{broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1, 2]> : tensor<3xi64>>}> : (tensor<1x32x256xf32>) -> tensor<4x32x256xf32>
  return %0 : tensor<4x32x256xf32>
}

//CHECK:func.func private @broadcast_in_dim(%arg0: tensor<1x32x256xf32>) -> tensor<4x32x256xf32> {