| 38 | namespace { |
| 39 | |
| 40 | Status GetStackShape(xla::XlaBuilder* builder, XlaResource* resource, |
| 41 | TensorShape* stack_shape) { |
| 42 | auto shape_or_status = builder->GetShape(resource->value()); |
| 43 | if (!shape_or_status.ok()) { |
| 44 | return shape_or_status.status(); |
| 45 | } |
| 46 | xla::Shape shape = shape_or_status.ValueOrDie(); |
| 47 | TF_RET_CHECK(shape.IsTuple()); |
| 48 | return XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape(shape, 0), |
| 49 | stack_shape); |
| 50 | } |
| 51 | |
| 52 | // Since the element shape is not provided to the Stack operator, |
| 53 | // we lazily initialize the Stack at the time of the first write. |