message: nhwc plumbing on pool grad op
diff: @@ -317,6 +317,7 @@ void DnnPoolingGradOp<T>::Compute(
  return;
  }
  +#if CUDNN_VERSION < 7300
  /// For now, cudnn does not support NHWC format, so we need to convert it
  /// to NCHW before calling cudnn. We need to get rid of this once it is done
  Tensor transformed_input;
  @@ -382,6 +383,39 @@ void DnnPoolingGradOp<T>::Compute(
  context->eigen_device<Device>(), out_backprop.tensor<T, 4>(),
  transformed_output_backprop.tensor<T, 4>());
  }
  +#else
  + Tensor transformed_input;
  + if (!tensor_in) {
  + OP_REQUIRES_OK(context, context->allocate_temp(DataTypeToEnum<T>::value,
  + tensor_in_shape,
  + &transformed_input));
  + } else {
  + transformed_input = *tensor_in;
  + }
  + Tensor transformed_output;
  + if (!tensor_out) {
  + OP_REQUIRES_OK(context, context->allocate_temp(DataTypeToEnum<T>::value,
  + out_backprop.shape(),
  + &transformed_output));
  + } else {
  + transformed_output = *tensor_out;
  + }
  + Tensor transformed_input_backprop = *input_backprop;
  + Tensor transformed_output_backprop = out_backprop;
  + se::dnn::DataLayout data_layout;
  + switch (data_format) {
  + case FORMAT_NHWC:
  + data_layout = se::dnn::DataLayout::kBatchYXDepth;
  + break;
  + case FORMAT_NCHW:
  + data_layout = se::dnn::DataLayout::kBatchDepthYX;
  + break;
  + default:
  + OP_REQUIRES(context, false,
  + errors::InvalidArgument("Unsupported format: ",
  + ToString(data_format)));
  + }
  +#endif // CUDNN_VERSION < 7300
  /// Get ready to call cudnn
  se::dnn::PoolingDescriptor pooling_desc;
  @@ -399,14 +433,14 @@ void DnnPoolingGradOp<T>::Compute(
  .set_height(params.out_height)
  .set_width(params.out_width)
  .set_feature_map_count(params.depth)
  - .set_layout(se::dnn::DataLayout::kBatchDepthYX);
  + .set_layout(data_layout);
  se::dnn::BatchDescriptor orig_input_desc;
  orig_input_desc.set_count(params.tensor_in_batch)
  .set_height(params.tensor_in_rows)
  .set_width(params.tensor_in_cols)
  .set_feature_map_count(params.depth)
  - .set_layout(se::dnn::DataLayout::kBatchDepthYX);
  + .set_layout(data_layout);
  auto orig_output_data =
  AsDeviceMemory(transformed_output.template flat<T>().data(),
  @@ -449,6 +483,7 @@ void DnnPoolingGradOp<T>::Compute(
  OP_REQUIRES(context, status,
  errors::Internal("dnn PoolBackward launch failed"));
  +#if CUDNN_VERSION < 7300
  if (data_format == FORMAT_NHWC) {
  /// Transform the output data from NCHW back to NHWC.
  auto toConstTensor = [](const Tensor& x) -> const Tensor { return x; };
  @@ -457,6 +492,7 @@ void DnnPoolingGradOp<T>::Compute(
  toConstTensor(transformed_input_backprop).template tensor<T, 4>(),
  input_backprop->tensor<T, 4>());
  }
  +#endif // CUDNN_VERSION < 7300
  }
  #define DEFINE_DNN_OPS(T) \
  