| 578 | } |
| 579 | |
| 580 | Status BiasAddGradShape(shape_inference::InferenceContext* c) { |
| 581 | ShapeHandle input_shape; |
| 582 | // Fetch the data_format attribute, which may not exist. |
| 583 | string data_format; |
| 584 | Status s = c->GetAttr("data_format", &data_format); |
| 585 | |
| 586 | if (s.ok() && data_format == "NCHW") { |
| 587 | TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &input_shape)); |
| 588 | c->set_output(0, c->Vector(c->Dim(input_shape, 1))); |
| 589 | } else { |
| 590 | TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape)); |
| 591 | c->set_output(0, c->Vector(c->Dim(input_shape, -1))); |
| 592 | } |
| 593 | |
| 594 | return Status::OK(); |
| 595 | } |
| 596 | |
| 597 | Status CheckFormatConstraintsOnShape(const TensorFormat tensor_format, |
| 598 | const ShapeHandle shape_handle, |