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

  %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> {