auto inputs = bn_op.getOperand();
    auto input_type = mlir::dyn_cast<RankedTensorType>(inputs.getType());
    if (!input_type) {
      return failure();
    }
    auto feature_index = bn_op.getFeatureIndex();

    // Compute mean
    int64_t input_last_dim = input_type.getRank() - 1;
    auto dims_type = RankedTensorType::get(/*shape=*/{input_last_dim},