(*result)[id].push_back(ones_like);
        }
      }
    } else {
      (*result)[id].push_back(output_gradients[i]);
    }
  }
  return absl::OkStatus();
}

// TODO(agarwal): use an automatic mechanism for handling None arguments to
// gradient functions.
//
// Some gradient functions can accept None arguments for gradients. The