| 220 | } |
| 221 | |
| 222 | Status MatMulShape(shape_inference::InferenceContext* c) { |
| 223 | ShapeHandle a; |
| 224 | TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &a)); |
| 225 | |
| 226 | ShapeHandle b; |
| 227 | TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &b)); |
| 228 | |
| 229 | bool transpose_a, transpose_b; |
| 230 | TF_RETURN_IF_ERROR(c->GetAttr("transpose_a", &transpose_a)); |
| 231 | TF_RETURN_IF_ERROR(c->GetAttr("transpose_b", &transpose_b)); |
| 232 | DimensionHandle output_rows = transpose_a ? c->Dim(a, 1) : c->Dim(a, 0); |
| 233 | DimensionHandle output_cols = transpose_b ? c->Dim(b, 0) : c->Dim(b, 1); |
| 234 | |
| 235 | // Validate that the inner shapes are compatible. |
| 236 | DimensionHandle inner_a = transpose_a ? c->Dim(a, 0) : c->Dim(a, 1); |
| 237 | DimensionHandle inner_b = transpose_b ? c->Dim(b, 1) : c->Dim(b, 0); |
| 238 | DimensionHandle merged; |
| 239 | TF_RETURN_IF_ERROR(c->Merge(inner_a, inner_b, &merged)); |
| 240 | |
| 241 | c->set_output(0, c->Matrix(output_rows, output_cols)); |
| 242 | return Status::OK(); |
| 243 | } |
| 244 | |
| 245 | Status MatMulGradFilterShape(shape_inference::InferenceContext* c) { |
| 246 | ShapeHandle a; |