| 519 | // -------------------------------------------------------------------------- |
| 520 | |
| 521 | Status BiasAddShape(shape_inference::InferenceContext* c) { |
| 522 | ShapeHandle input_shape; |
| 523 | |
| 524 | // Fetch the data_format attribute, which may not exist. |
| 525 | string data_format; |
| 526 | Status s = c->GetAttr("data_format", &data_format); |
| 527 | |
| 528 | if (s.ok() && data_format == "NCHW") { |
| 529 | TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &input_shape)); |
| 530 | } else { |
| 531 | TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape)); |
| 532 | } |
| 533 | |
| 534 | ShapeHandle bias_shape; |
| 535 | TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &bias_shape)); |
| 536 | DimensionHandle bias_dim = c->Dim(bias_shape, 0); |
| 537 | |
| 538 | // If rank unknown, return unknown shape. |
| 539 | if (!c->RankKnown(input_shape)) { |
| 540 | c->set_output(0, c->UnknownShape()); |
| 541 | return Status::OK(); |
| 542 | } |
| 543 | |
| 544 | // Output has the same shape as the input, and matches the length of |
| 545 | // the bias in its bias dimension. |
| 546 | ShapeHandle output_shape; |
| 547 | if (s.ok() && data_format == "NCHW") { |
| 548 | // Merge the length of bias_shape into the third to last dimension |
| 549 | ShapeHandle first; |
| 550 | TF_RETURN_IF_ERROR(c->Subshape(input_shape, 0, 1, &first)); |
| 551 | |
| 552 | ShapeHandle last; |
| 553 | TF_RETURN_IF_ERROR(c->Subshape(input_shape, 2, &last)); |
| 554 | |
| 555 | DimensionHandle input_bias_dim = c->Dim(input_shape, 1); |
| 556 | DimensionHandle merged_bias_dim; |
| 557 | TF_RETURN_IF_ERROR(c->Merge(input_bias_dim, bias_dim, &merged_bias_dim)); |
| 558 | ShapeHandle merged_bias = c->Vector(merged_bias_dim); |
| 559 | |
| 560 | ShapeHandle temp; |
| 561 | TF_RETURN_IF_ERROR(c->Concatenate(first, merged_bias, &temp)); |
| 562 | TF_RETURN_IF_ERROR(c->Concatenate(temp, last, &output_shape)); |
| 563 | } else { |
| 564 | ShapeHandle all_but_bias; |
| 565 | TF_RETURN_IF_ERROR(c->Subshape(input_shape, 0, -1, &all_but_bias)); |
| 566 | |
| 567 | DimensionHandle input_bias_dim = c->Dim(input_shape, -1); |
| 568 | DimensionHandle merged_bias_dim; |
| 569 | TF_RETURN_IF_ERROR(c->Merge(input_bias_dim, bias_dim, &merged_bias_dim)); |
| 570 | |
| 571 | ShapeHandle merged_bias = c->Vector(merged_bias_dim); |
| 572 | TF_RETURN_IF_ERROR( |
| 573 | c->Concatenate(all_but_bias, merged_bias, &output_shape)); |
| 574 | } |
| 575 | |
| 576 | c->set_output(0, output_shape); |
| 577 | return Status::OK(); |
| 578 | } |