Since the element shape is not provided to the Stack operator, we lazily initialize the Stack at the time of the first write. If a Stack `resource` has not been initialized, constructs storage for the Stack with elements of `elem_shape`. For both initialized and uninitialized Stacks, checks that the tensor has a type compatible with 'dtype' and shape compatible with 'elem_shape'. TODO(phawkins):
| 60 | // TODO(phawkins): consider changing the API of the stack operators to |
| 61 | // allow an optional element shape at stack construction time. |
| 62 | Status MaybeInitializeStack(xla::XlaBuilder* builder, XlaResource* resource, |
| 63 | DataType dtype, const TensorShape& elem_shape) { |
| 64 | if (resource->type() != dtype) { |
| 65 | return errors::InvalidArgument( |
| 66 | "Stack dtype is ", DataTypeString(resource->type()), |
| 67 | " but op has dtype ", DataTypeString(dtype), "."); |
| 68 | } |
| 69 | |
| 70 | TensorShape stack_shape; |
| 71 | stack_shape.AddDim(resource->max_array_size()); |
| 72 | stack_shape.AppendShape(elem_shape); |
| 73 | |
| 74 | if (!resource->initialized()) { |
| 75 | // Stack has not been initialized. |
| 76 | TF_RETURN_IF_ERROR(resource->SetTypeAndShape(dtype, elem_shape)); |
| 77 | TF_RETURN_IF_ERROR(resource->SetZeroValue(builder)); |
| 78 | } else { |
| 79 | // Checks the expected shape matches the actual shape. |
| 80 | TensorShape actual_shape; |
| 81 | TF_RETURN_IF_ERROR(GetStackShape(builder, resource, &actual_shape)); |
| 82 | if (stack_shape != actual_shape) { |
| 83 | return errors::InvalidArgument( |
| 84 | "Mismatched Stack shapes: ", stack_shape.DebugString(), " vs ", |
| 85 | actual_shape.DebugString()); |
| 86 | } |
| 87 | } |
| 88 | return Status::OK(); |
| 89 | } |
| 90 | |
| 91 | class StackOp : public XlaOpKernel { |
| 92 | public: |
no test coverage detected