Method AllocateOutputTensor

tensorflow/core/kernels/mkl_slice_op.cc:485–516

 private:
  void AllocateOutputTensor(OpKernelContext* context,
                            const MklDnnShape& input_mkl_shape,
                            MEMORY_PRIMITIVE_DESC* output_pd,
                            const memory::dims& output_dims,
                            Tensor** output_tensor,
                            MklDnnShape* output_mkl_shape) {
    DCHECK(output_tensor);
    DCHECK(output_mkl_shape);

    TensorShape output_tf_shape;

    if (input_mkl_shape.IsMklTensor()) {
      // Since input tensor is in OneDNN layout, output tensor will be in OneDNN
      // layout.

      // Allocate shape of OneDNN tensor.
      output_mkl_shape->SetMklTensor(true);
      output_mkl_shape->SetMklLayout(output_pd);
      output_mkl_shape->SetElemType(MklDnnType<T>());
      output_mkl_shape->SetTfLayout(input_mkl_shape.GetDimension(), output_dims,
                                    input_mkl_shape.GetTfDataFormat());

      output_tf_shape.AddDim(output_pd->get_size() / sizeof(T));
    } else {
      // If input is not in OneDNN layout, then output won't be in OneDNN layout.
      output_mkl_shape->SetMklTensor(false);
      output_tf_shape = MklDnnDimsToTFShape(output_dims);
    }

    AllocateOutputSetMklShape(context, 0, output_tensor, output_tf_shape,
                              *output_mkl_shape);
  }
};

// OneDNN Slice registration