Searched refs:out_rank (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | reduction_gpu_kernels.cu.h | 1074 int in_dim0, int in_dim1, int in_dim2, int out_rank,
1078 if (out_rank == 0) {
1081 } else if (in_rank == 2 && out_rank == 1 &&
1084 } else if (in_rank == 2 && out_rank == 1 &&
1087 } else if (in_rank == 3 && out_rank == 2 && reduction_axes[0] == 1) {
1096 } else if (in_rank == 3 && out_rank == 1 && reduction_axes[0] == 0 &&
1103 << " " << out_rank;
1104 if (out_rank == 1) ss << " " << reduction_axes[0];
1105 if (out_rank == 2) ss << " " << reduction_axes[1];
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | array_ops.cc | 492 int out_rank = c->Rank(out); in __anondb9326b20702() local
493 TF_RETURN_IF_ERROR(c->WithRankAtMost(in, out_rank, &in)); in __anondb9326b20702()
500 auto out_dim = c->Dim(out, out_rank - i - 1); in __anondb9326b20702()
503 c->ReplaceDim(out, out_rank - i - 1, out_dim, &out)); in __anondb9326b20702()
|