| 80 | } |
| 81 | |
| 82 | Status PadShapeFn(InferenceContext* c) { |
| 83 | // Paddings is a matrix of [input_rank, 2]. |
| 84 | ShapeHandle paddings; |
| 85 | TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &paddings)); |
| 86 | DimensionHandle unused; |
| 87 | TF_RETURN_IF_ERROR(c->WithValue(c->Dim(paddings, 1), 2, &unused)); |
| 88 | |
| 89 | // n_dim and input.rank are equivalent. |
| 90 | ShapeHandle input = c->input(0); |
| 91 | DimensionHandle n_dim = c->Dim(paddings, 0); |
| 92 | if (c->ValueKnown(n_dim)) { |
| 93 | TF_RETURN_IF_ERROR(c->WithRank(input, c->Value(n_dim), &input)); |
| 94 | } else if (c->RankKnown(input)) { |
| 95 | TF_RETURN_IF_ERROR(c->WithValue(n_dim, c->Rank(input), &n_dim)); |
| 96 | } |
| 97 | |
| 98 | const Tensor* paddings_t = c->input_tensor(1); |
| 99 | |
| 100 | // paddings_t is unknown |
| 101 | if (paddings_t == nullptr) { |
| 102 | if (c->ValueKnown(n_dim)) { |
| 103 | // Make output with n_dim unknown dims. |
| 104 | c->set_output(0, c->UnknownShapeOfRank(c->Value(n_dim))); |
| 105 | } else { |
| 106 | c->set_output(0, c->UnknownShape()); |
| 107 | } |
| 108 | return Status::OK(); |
| 109 | } |
| 110 | |
| 111 | const int64 num_dims = paddings_t->shape().dim_size(0); |
| 112 | TF_RETURN_IF_ERROR(c->WithRank(input, num_dims, &input)); |
| 113 | TF_RETURN_IF_ERROR(c->WithValue(n_dim, num_dims, &n_dim)); |
| 114 | |
| 115 | if (paddings_t->dtype() == DT_INT32) { |
| 116 | return PadKnown<int32>(c, input, paddings_t, num_dims); |
| 117 | } else { |
| 118 | return PadKnown<int64>(c, input, paddings_t, num_dims); |
| 119 | } |
| 120 | } |
| 121 | |
| 122 | Status TransposeShapeFn(InferenceContext* c) { |
| 123 | ShapeHandle input = c->input(0); |
nothing calls this directly
no test coverage detected