int M = f.getNumArguments() - N;
  DCOMMENT_OP(op, "Inferring shape for with N = " << N << " and M = " << M);

  // Initialize with function input types.
  auto input_types = llvm::to_vector<1>(
      cast<FunctionOpInterface>(f.getOperation()).getArgumentTypes());

  DatasetInput input_elements = GetDatasetInput(op->getOperand(0));
  if (!input_elements) {
    op->emitWarning("unexpected dataset input; skipping function refinement");