| 120 | } |
| 121 | |
| 122 | Status TransposeShapeFn(InferenceContext* c) { |
| 123 | ShapeHandle input = c->input(0); |
| 124 | ShapeHandle perm_shape = c->input(1); |
| 125 | const Tensor* perm = c->input_tensor(1); |
| 126 | DimensionHandle perm_elems = c->NumElements(perm_shape); |
| 127 | // If we don't have rank information on the input or value information on |
| 128 | // perm we can't return any shape information, otherwise we have enough |
| 129 | // information to at least find the rank of the output. |
| 130 | if (!c->RankKnown(input) && !c->ValueKnown(perm_elems) && perm == nullptr) { |
| 131 | c->set_output(0, c->UnknownShape()); |
| 132 | return Status::OK(); |
| 133 | } |
| 134 | |
| 135 | // Find our value of the rank. |
| 136 | int64 rank; |
| 137 | if (c->RankKnown(input)) { |
| 138 | rank = c->Rank(input); |
| 139 | } else if (c->ValueKnown(perm_elems)) { |
| 140 | rank = c->Value(perm_elems); |
| 141 | } else { |
| 142 | rank = perm->NumElements(); |
| 143 | } |
| 144 | if (!c->RankKnown(input) && rank < 2) { |
| 145 | // A permutation array containing a single element is ambiguous. It could |
| 146 | // indicate either a scalar or a 1-dimensional array, both of which the |
| 147 | // transpose op returns unchanged. |
| 148 | c->set_output(0, input); |
| 149 | return Status::OK(); |
| 150 | } |
| 151 | |
| 152 | std::vector<DimensionHandle> dims; |
| 153 | dims.resize(rank); |
| 154 | TF_RETURN_IF_ERROR(c->WithRank(input, rank, &input)); |
| 155 | // Ensure that perm is a vector and has rank elements. |
| 156 | TF_RETURN_IF_ERROR(c->WithRank(perm_shape, 1, &perm_shape)); |
| 157 | TF_RETURN_IF_ERROR(c->WithValue(perm_elems, rank, &perm_elems)); |
| 158 | |
| 159 | // If we know the rank of the input and the value of perm, we can return |
| 160 | // all shape informantion, otherwise we can only return rank information, |
| 161 | // but no information for the dimensions. |
| 162 | if (perm != nullptr) { |
| 163 | std::vector<int64> data; |
| 164 | if (perm->dtype() == DT_INT32) { |
| 165 | data = AsInt64<int32>(perm, rank); |
| 166 | } else { |
| 167 | data = AsInt64<int64>(perm, rank); |
| 168 | } |
| 169 | |
| 170 | for (int32 i = 0; i < rank; ++i) { |
| 171 | int64 in_idx = data[i]; |
| 172 | if (in_idx >= rank || in_idx <= -rank) { |
| 173 | return errors::InvalidArgument("perm dim ", in_idx, |
| 174 | " is out of range of input rank ", rank); |
| 175 | } |
| 176 | dims[i] = c->Dim(input, in_idx); |
| 177 | } |
| 178 | } else { |
| 179 | for (int i = 0; i < rank; ++i) { |
nothing calls this directly
no test coverage detected