| 96 | } |
| 97 | |
| 98 | void Compile(XlaOpKernelContext* ctx) override { |
| 99 | int64 max_size; |
| 100 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &max_size)); |
| 101 | OP_REQUIRES( |
| 102 | ctx, max_size >= 0, |
| 103 | errors::InvalidArgument( |
| 104 | "XLA compilation requires a fixed stack size upper bound. If " |
| 105 | "you are using tf.while_loop, set the maximum_iterations parameter " |
| 106 | "to fix this issue.")); |
| 107 | |
| 108 | // We defer initializing the Stack resource until we see the first push. |
| 109 | // Otherwise we do not know the shape of the stack elements. |
| 110 | XlaResource* resource = |
| 111 | ctx->xla_context()->AddResource(XlaResource::CreateStack( |
| 112 | /*name=*/absl::StrCat("Stack: ", stack_name_), dtype_, max_size)); |
| 113 | ctx->SetResourceOutput(0, resource); |
| 114 | } |
| 115 | |
| 116 | private: |
| 117 | DataType dtype_; |
nothing calls this directly
no test coverage detected