  void Compile(XlaOpKernelContext* ctx) override {
    xla::Literal paddings;
    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(1, &paddings));

    SpaceToBatch(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0),
                 {block_size_, block_size_}, paddings);
  }
  void Compile(XlaOpKernelContext* ctx) override {
    xla::Literal paddings;
    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(1, &paddings));

    SpaceToBatch(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0),
                 {block_size_, block_size_}, paddings);
  }