| 106 | } // namespace |
| 107 | |
| 108 | Status MklTransposeCpuOp::DoTranspose(OpKernelContext* ctx, const Tensor& in, |
| 109 | gtl::ArraySlice<int32> perm, |
| 110 | Tensor* out) { |
| 111 | // OneDNN has limit on the maximum number of dimensions in a tensor. |
| 112 | // Fallback to Eigen for not supported cases. |
| 113 | if (in.dims() <= TENSOR_MAX_DIMS) { |
| 114 | switch (in.dtype()) { |
| 115 | case DT_FLOAT: |
| 116 | return MKLTransposeND<float>(ctx, in, out, perm); |
| 117 | break; |
| 118 | case DT_BFLOAT16: |
| 119 | return MKLTransposeND<bfloat16>(ctx, in, out, perm); |
| 120 | break; |
| 121 | // TODO(nhasabni): support other types such as INT8. |
| 122 | default: |
| 123 | break; |
| 124 | } |
| 125 | } |
| 126 | |
| 127 | // Fallback to eigen if transpose parameters not supported by OneDNN |
| 128 | typedef Eigen::ThreadPoolDevice CPUDevice; |
| 129 | return ::tensorflow::DoTranspose(ctx->eigen_device<CPUDevice>(), in, perm, |
| 130 | out); |
| 131 | } |
| 132 | |
| 133 | Status MklConjugateTransposeCpuOp::DoTranspose(OpKernelContext* ctx, |
| 134 | const Tensor& in, |
nothing calls this directly
no test coverage detected