Method Compile

tensorflow/compiler/tf2xla/kernels/shape_op.cc:244–277

  explicit ZerosLikeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}

  void Compile(XlaOpKernelContext* ctx) override {
    if (IsTensorListInput(ctx, 0)) {
      // Input is a TensorList.

      // Check the TensorList input is initialized.
      xla::XlaOp list = ctx->Input(0);
      bool is_initialized;
      OP_REQUIRES_OK(ctx, IsTensorListInitialized(list, &is_initialized));
      OP_REQUIRES(
          ctx, is_initialized,
          errors::InvalidArgument(
              "TensorList input for ZerosLike op is an uninitialized list"));

      auto list_shape_or = ctx->builder()->GetShape(list);
      OP_REQUIRES_OK(ctx, list_shape_or.status());
      xla::XlaOp new_list;
      OP_REQUIRES_OK(
          ctx, CreateZerosTensorListWithShape(
                   ctx->builder(), list_shape_or.ValueOrDie(), &new_list));

      xla::XlaOp push_index;
      OP_REQUIRES_OK(ctx, GetTensorListPushIndex(list, &push_index));

      xla::XlaOp result;
      OP_REQUIRES_OK(ctx,
                     SetTensorListPushIndex(new_list, push_index, &result));
      ctx->SetTensorListOutput(0, result);
    } else {
      const TensorShape input_shape = ctx->InputShape(0);

      auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0));
      ctx->SetOutput(0, xla::Broadcast(zero, input_shape.dim_sizes()));
    }
  }
};

REGISTER_XLA_OP(Name("ZerosLike").AllowVariantTypes(), ZerosLikeOp);