| 191 | |
| 192 | private: |
| 193 | void AllocateOutputTensor( |
| 194 | OpKernelContext* context, |
| 195 | const lrn_forward::primitive_desc& lrn_fwd_prim_desc, |
| 196 | const memory::dims output_dims_mkl_order, |
| 197 | const MKL_TENSOR_FORMAT& output_tf_format, Tensor** output_tensor) { |
| 198 | DCHECK(output_tensor != nullptr); |
| 199 | MEMORY_PRIMITIVE_DESC dst_pd = lrn_fwd_prim_desc.PRIMITIVE_DESC_DST; |
| 200 | |
| 201 | MklDnnShape output_mkl_shape; |
| 202 | // We only handle the case when the inputs and output are in OneDNN format |
| 203 | // Any other case is handled by Eigen |
| 204 | output_mkl_shape.SetMklTensor(true); |
| 205 | output_mkl_shape.SetMklLayout(&dst_pd); |
| 206 | output_mkl_shape.SetElemType(MklDnnType<T>()); |
| 207 | output_mkl_shape.SetTfLayout(output_dims_mkl_order.size(), |
| 208 | output_dims_mkl_order, output_tf_format); |
| 209 | TensorShape output_tf_shape; |
| 210 | // only allocate enough space for the elements we need. |
| 211 | size_t num_bytes = dst_pd.get_size(); |
| 212 | CHECK_EQ(num_bytes % sizeof(T), 0); |
| 213 | output_tf_shape.AddDim(num_bytes / sizeof(T)); |
| 214 | AllocateOutputSetMklShape(context, kIdxOutput, output_tensor, |
| 215 | output_tf_shape, output_mkl_shape); |
| 216 | } |
| 217 | |
| 218 | // Fallback implementation - Taken from lrn_op.cc |
| 219 | // TODO(inteltf) Check if we can use EigenLRNOp directly instead of making a |
no test coverage detected