| 456 | } |
| 457 | |
| 458 | void AllocateOutputTensor( |
| 459 | OpKernelContext* context, |
| 460 | const lrn_backward::primitive_desc& lrn_bkwd_prim_desc, |
| 461 | const memory::dims output_dims_mkl_order, |
| 462 | const MKL_TENSOR_FORMAT& output_tf_format, Tensor** output_tensor) { |
| 463 | DCHECK(output_tensor != nullptr); |
| 464 | MEMORY_PRIMITIVE_DESC dst_pd = lrn_bkwd_prim_desc.PRIMITIVE_DESC_DIFF_SRC; |
| 465 | MklDnnShape output_mkl_shape; |
| 466 | |
| 467 | // We assume that all outputs at this point are OneDNN Tensors |
| 468 | output_mkl_shape.SetMklTensor(true); |
| 469 | output_mkl_shape.SetMklLayout(&dst_pd); |
| 470 | output_mkl_shape.SetElemType(MklDnnType<T>()); |
| 471 | output_mkl_shape.SetTfLayout(output_dims_mkl_order.size(), |
| 472 | output_dims_mkl_order, output_tf_format); |
| 473 | |
| 474 | TensorShape output_tf_shape; |
| 475 | size_t num_bytes = dst_pd.get_size(); |
| 476 | CHECK_EQ(num_bytes % sizeof(T), 0); |
| 477 | output_tf_shape.AddDim(num_bytes / sizeof(T)); |
| 478 | AllocateOutputSetMklShape(context, kIdxOutput, output_tensor, |
| 479 | output_tf_shape, output_mkl_shape); |
| 480 | } |
| 481 | |
| 482 | memory::desc ConfigureInputGradient(const Tensor& input_grad_tensor, |
| 483 | const MklDnnShape& input_grad_dnn_shape, |
nothing calls this directly
no test coverage detected