  void Compute(OpKernelContext* ctx) override {
    const Tensor& shape_t = ctx->input(0);
    const Tensor& rate_t = ctx->input(1);

    TensorShape samples_shape;
    OP_REQUIRES_OK(ctx, tensor::MakeShape(shape_t, &samples_shape));
    const int64 num_samples = samples_shape.num_elements();

    samples_shape.AppendShape(rate_t.shape());
    // Allocate output samples.
    Tensor* samples_t = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, samples_shape, &samples_t));
    if (num_samples == 0) return;

    const auto rate_flat = rate_t.flat<T>().data();
    const int64 num_rate = rate_t.NumElements();
    auto samples_flat = samples_t->flat<U>().data();
    random::PhiloxRandom rng = generator_.ReserveRandomOutputs(
        num_samples * num_rate, kReservedSamplesPerOutput);

    functor::PoissonFunctor<CPUDevice, T, U>()(
        ctx, ctx->eigen_device<CPUDevice>(), rate_flat, num_rate, num_samples,
        rng, samples_flat);
  }
  void Compute(OpKernelContext* ctx) override {
    const Tensor& shape_t = ctx->input(0);
    const Tensor& rate_t = ctx->input(1);

    TensorShape samples_shape;
    OP_REQUIRES_OK(ctx, tensor::MakeShape(shape_t, &samples_shape));
    const int64 num_samples = samples_shape.num_elements();

    samples_shape.AppendShape(rate_t.shape());
     
    Tensor* samples_t = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, samples_shape, &samples_t));
    if (num_samples == 0) return;

    const auto rate_flat = rate_t.flat<T>().data();
    const int64 num_rate = rate_t.NumElements();
    auto samples_flat = samples_t->flat<U>().data();
    random::PhiloxRandom rng = generator_.ReserveRandomOutputs(
        num_samples * num_rate, kReservedSamplesPerOutput);

    functor::PoissonFunctor<CPUDevice, T, U>()(
        ctx, ctx->eigen_device<CPUDevice>(), rate_flat, num_rate, num_samples,
        rng, samples_flat);
  }