| 102 | } |
| 103 | |
| 104 | xla::StatusOr<absl::optional<Tensor>> XlaExpression::ResolveConstant( |
| 105 | xla::Client* client, bool dynamic_dimension_is_minus_one) const { |
| 106 | switch (kind()) { |
| 107 | case Kind::kConstant: |
| 108 | return {constant_value()}; |
| 109 | case Kind::kXlaOp: |
| 110 | break; |
| 111 | case Kind::kTensorList: |
| 112 | TF_FALLTHROUGH_INTENDED; |
| 113 | case Kind::kResource: |
| 114 | TF_FALLTHROUGH_INTENDED; |
| 115 | case Kind::kInvalid: |
| 116 | return errors::InvalidArgument( |
| 117 | "ResolveConstant called on XlaExpression: ", HumanString()); |
| 118 | } |
| 119 | |
| 120 | TF_ASSIGN_OR_RETURN(bool is_constant, |
| 121 | handle().builder()->IsConstant(handle())); |
| 122 | if (!is_constant) return {absl::nullopt}; |
| 123 | |
| 124 | TF_ASSIGN_OR_RETURN(xla::XlaComputation constant_graph, |
| 125 | handle().builder()->BuildConstantSubGraph( |
| 126 | handle(), dynamic_dimension_is_minus_one)); |
| 127 | |
| 128 | TF_ASSIGN_OR_RETURN(TensorShape shape, GetShape()); |
| 129 | |
| 130 | // The XLA layout is specified minor to major, and TensorFlow uses a major to |
| 131 | // minor order. |
| 132 | std::vector<int64> layout_indices(shape.dims()); |
| 133 | std::iota(layout_indices.rbegin(), layout_indices.rend(), 0); |
| 134 | xla::Layout layout = xla::LayoutUtil::MakeLayout(layout_indices); |
| 135 | TF_ASSIGN_OR_RETURN(xla::Literal literal, |
| 136 | client->ComputeConstant(constant_graph, &layout)); |
| 137 | Tensor tensor; |
| 138 | TF_RETURN_IF_ERROR(LiteralToHostTensor(literal, dtype(), &tensor)); |
| 139 | return {tensor}; |
| 140 | } |
| 141 | |
| 142 | xla::StatusOr<TensorShape> XlaExpression::GetShape() const { |
| 143 | switch (kind_) { |