| 140 | } |
| 141 | |
| 142 | xla::StatusOr<TensorShape> XlaExpression::GetShape() const { |
| 143 | switch (kind_) { |
| 144 | case Kind::kConstant: |
| 145 | return constant_value().shape(); |
| 146 | case Kind::kXlaOp: { |
| 147 | TF_ASSIGN_OR_RETURN(xla::Shape xla_shape, |
| 148 | handle().builder()->GetShape(handle())); |
| 149 | TensorShape shape; |
| 150 | TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, &shape)); |
| 151 | return shape; |
| 152 | } |
| 153 | case Kind::kTensorList: |
| 154 | return TensorShape({}); |
| 155 | case Kind::kResource: |
| 156 | return TensorShape({}); |
| 157 | case Kind::kInvalid: |
| 158 | return errors::InvalidArgument( |
| 159 | "GetShape() called on invalid XlaExpression"); |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | } // namespace tensorflow |