const mlir::ArrayAttr partition_dims = op.getPartitionDims();
  for (const mlir::Attribute &dim : partition_dims) {
    num_partitions *= dim.cast<IntegerAttr>().getInt();
  }

  const bool is_packed = op.getIsPacked();
  const bool replicated = partition_dims.empty();
  const int num_inputs_expected = is_packed ? 1 : num_partitions;

  if (!((replicated && !is_packed) || (op.getN() == num_inputs_expected))) {