Method Compute

tensorflow/core/kernels/stack.cc:164–210

}

void StackOp::Compute(OpKernelContext* ctx) {
  int32 size = std::numeric_limits<int32>::max();
  if (ctx->num_inputs() > 0) {
    const Tensor* tensor_size;
    OP_REQUIRES_OK(ctx, ctx->input("max_size", &tensor_size));

    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(tensor_size->shape()),
        errors::InvalidArgument("Stack size must be a scalar, but had shape: ",
                                tensor_size->shape().DebugString()));

    int32 size_value = tensor_size->scalar<int32>()();
    if (size_value >= 0) {
      size = size_value;
    }
  }

  static const char kContainer[] = "_stacks";
  auto stack_id = Stack::stack_counter.fetch_add(1);
  string stack_name = strings::StrCat(stack_name_, "_", stack_id);
  // Store the handle in a per-step container.
  ResourceMgr* rm = ctx->resource_manager();
  OP_REQUIRES(ctx, rm != nullptr, errors::Internal("No resource manager."));
  string key = strings::StrCat(kContainer, stack_name);
  auto* step_container = ctx->step_container();
  OP_REQUIRES(ctx, step_container != nullptr,
              errors::Internal("No step container."));
  Stack* stack = new Stack(elem_type_, stack_name, size);
  OP_REQUIRES_OK(ctx, step_container->Create(rm, key, stack));
  if (IsRefType(ctx->expected_output_dtype(0))) {
    // Create the stack handle.
    AllocatorAttributes alloc_attr;
    alloc_attr.set_on_host(true);
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(tensorflow::DT_STRING,
                                           tensorflow::TensorShape({2}),
                                           &stack->handle_, alloc_attr));
    auto handle = stack->handle_.flat<tstring>();
    handle(0) = kContainer;
    handle(1) = std::move(stack_name);
    ctx->set_output_ref(0, stack->mu(), &stack->handle_);
  } else {
    Tensor* handle;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &handle));
    handle->flat<ResourceHandle>()(0) =
        ctx->step_container()->MakeResourceHandle<Stack>(key, *ctx->device());
  }
}

// StackPushOp