Method Compile

tensorflow/compiler/tf2xla/kernels/stack_ops.cc:133–164

  }

  void Compile(XlaOpKernelContext* ctx) override {
    xla::XlaBuilder* b = ctx->builder();
    TensorShape elem_shape = ctx->InputShape(1);

    XlaResource* resource;
    OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));

    // Initializes the Stack, if the element shape was not already known.
    OP_REQUIRES_OK(ctx, MaybeInitializeStack(b, resource, dtype_, elem_shape));

    xla::XlaOp ta = xla::GetTupleElement(resource->value(), 0);
    xla::XlaOp index = xla::GetTupleElement(resource->value(), 1);
    xla::XlaOp value = ctx->Input(1);

    // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0].
    std::vector<xla::XlaOp> start_indices(elem_shape.dims() + 1,
                                          xla::ConstantR0<int32>(b, 0));
    start_indices[0] = index;

    TensorShape slice_shape = elem_shape;
    slice_shape.InsertDim(0, 1LL);
    auto update = xla::Reshape(value, slice_shape.dim_sizes());

    // TODO(phawkins): We don't check the index is in bounds --- there is no
    // error mechanism in XLA.
    OP_REQUIRES_OK(ctx,
                   resource->SetValue(xla::Tuple(
                       b, {xla::DynamicUpdateSlice(ta, update, start_indices),
                           xla::Add(index, xla::ConstantR0<int32>(b, 1))})));

    ctx->SetOutput(0, value);
  }

 private:
  DataType dtype_;