Function GetTfShape

tensorflow/core/util/mkl_util.h:729–743

Get shape of input tensor pointed by 'input_idx' in TensorShape format. If the input tensor is in OneDNN layout, then obtains TensorShape from MklShape.

/// If the input tensor is in OneDNN layout, then obtains TensorShape from
/// MklShape.
inline TensorShape GetTfShape(OpKernelContext* context, size_t input_idx,
                              bool eager_mode = false) {
  // Sanity check.
  CHECK_NOTNULL(context);
  CHECK_LT(input_idx, context->num_inputs());

  MklDnnShape input_mkl_shape;
  GetMklShape(context, input_idx, &input_mkl_shape, eager_mode);
  if (input_mkl_shape.IsMklTensor() && !eager_mode) {
    return input_mkl_shape.GetTfShape();
  } else {
    const Tensor& t = MklGetInput(context, input_idx);
    return t.shape();
  }
}

// Allocate the second output tensor that will contain
// the OneDNN shape serialized