Searched refs:out_rank (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | reduction_gpu_kernels.cu.h | 854 int in_dim0, int in_dim1, int in_dim2, int out_rank,
858 if (out_rank == 0) {
861 } else if (in_rank == 2 && out_rank == 1 &&
864 } else if (in_rank == 2 && out_rank == 1 &&
867 } else if (in_rank == 3 && out_rank == 2 && reduction_axes[0] == 1) {
870 } else if (in_rank == 3 && out_rank == 1 && reduction_axes[0] == 0 &&
877 << " " << out_rank;
878 if (out_rank == 1) ss << " " << reduction_axes[0];
879 if (out_rank == 2) ss << " " << reduction_axes[1];
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | array_ops.cc | 486 int out_rank = c->Rank(out); in __anondb9326b20702() local
487 TF_RETURN_IF_ERROR(c->WithRankAtMost(in, out_rank, &in)); in __anondb9326b20702()
494 auto out_dim = c->Dim(out, out_rank - i - 1); in __anondb9326b20702()
497 c->ReplaceDim(out, out_rank - i - 1, out_dim, &out)); in __anondb9326b20702()
|