A XlaExpression represents a symbolic TensorFlow value in a TF->XLA compilation. An expression is one of: a constant tensor. an xla::XlaOp, representing a symbolic XLA value. a resource, e.g., a variable, represented as an XlaResource pointer. a tensor list, represented by a tuple of tensors and the list length. Constant tensors are mostly an optimization to avoid passing large constants to XLA,
| 43 | // specific logic around shape management since the tuples are not supported by |
| 44 | // TensorFlow. |
| 45 | class XlaExpression { |
| 46 | public: |
| 47 | enum class Kind { |
| 48 | kInvalid, |
| 49 | kConstant, |
| 50 | kXlaOp, |
| 51 | kResource, |
| 52 | kTensorList, |
| 53 | }; |
| 54 | |
| 55 | XlaExpression(); |
| 56 | XlaExpression(const XlaExpression&) = default; |
| 57 | XlaExpression& operator=(const XlaExpression&) = default; |
| 58 | |
| 59 | // Builds an invalid expression. (Same as the default constructor, but makes |
| 60 | // the intent clearer.) |
| 61 | static XlaExpression Invalid(); |
| 62 | |
| 63 | // Builds a constant XLA expression. |
| 64 | static XlaExpression Constant(Tensor value); |
| 65 | |
| 66 | // Builds a XlaOp expression. Since the mapping from TF data types to XLA |
| 67 | // types is not 1-1, the TF type must also be provided; in general it cannot |
| 68 | // be derived from the XLA type. |
| 69 | static XlaExpression XlaOp(xla::XlaOp value, DataType dtype); |
| 70 | |
| 71 | // Builds a tensor list expression. |
| 72 | static XlaExpression TensorList(xla::XlaOp tensor_list); |
| 73 | |
| 74 | // Builds a resource expression. |
| 75 | static XlaExpression Resource(XlaResource* resource); |
| 76 | |
| 77 | Kind kind() const { return kind_; } |
| 78 | |
| 79 | DataType dtype() const { return dtype_; } |
| 80 | |
| 81 | // handle() returns the XlaOp that backs a kXlaOp expression. |
| 82 | const xla::XlaOp& handle() const { return handle_; } |
| 83 | |
| 84 | const Tensor& constant_value() const { return constant_value_; } |
| 85 | |
| 86 | XlaResource* resource() const { return resource_; } |
| 87 | |
| 88 | // Returns a human-readable summary of the expression. |
| 89 | string HumanString() const; |
| 90 | |
| 91 | // Returns the value of a kConstant or kXlaOp as an xla::XlaOp. Returns |
| 92 | // an erroneous XlaOp if the expression is not a constant or an expression. |
| 93 | xla::XlaOp AsXlaOp(xla::XlaBuilder* builder) const; |
| 94 | |
| 95 | // If a kXlaOp or kConstant expression can be resolved to a compile-time |
| 96 | // constant, returns the value as a host-memory Tensor. Returns an empty |
| 97 | // optional if it cannot be resolved. Returns an error if passed a resource |
| 98 | // expression. |
| 99 | xla::StatusOr<absl::optional<Tensor>> ResolveConstant( |
| 100 | xla::Client* client, bool dynamic_dimension_is_minus_one = false) const; |
| 101 | |
| 102 | // Returns the shape of the tensor. |
no outgoing calls