void operator()(OpKernelContext* ctx, const CPUDevice& d, int64 num_batches,
                  int64 samples_per_batch, int64 num_elements,
                  const BCastList<4>& bcast,
                  typename TTypes<T>::ConstFlat means,
                  typename TTypes<T>::ConstFlat stddevs,
                  typename TTypes<T>::ConstFlat minvals,
                  typename TTypes<T>::ConstFlat maxvals,
                  const random::PhiloxRandom& gen,
                  typename TTypes<T>::Flat output) {
    // The randn rejection sampling is used when the mean and at least this many
    // standard deviations are inside the bounds.
    // The uniform proposal samplers become less efficient as the bounds are
    // further from the mean, the reverse is true for the randn sampler.
    // This number was chosen by empirical benchmarking. If modified, the
    // benchmarks in parameterized_truncated_normal_op_test should also be
    // changed.
    const T kStdDevsInsideBoundsToUseRandnSampler = T(1.3);
    auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());

    auto do_work = [num_batches, samples_per_batch, &ctx, &bcast, &means,
                    &stddevs, &minvals, &maxvals, &gen, &output,
                    kStdDevsInsideBoundsToUseRandnSampler](int64 start_output,
                                                           int64 limit_output) {
      // Capturing "gen" by-value would only make a copy for the _shared_
      // lambda.  Since we want to let each worker have its own copy, we pass
      // "gen" by reference and explicitly do a copy assignment here.
      random::PhiloxRandom gen_copy = gen;
      using Uniform = random::UniformDistribution<random::PhiloxRandom, T>;
      Uniform dist;
      using Normal = random::NormalDistribution<random::PhiloxRandom, T>;
      Normal normal_dist;
      // Skip takes units of 128 bits. The Uniform::kResultElementCount - 1
      // is so rounding doesn't lead to
      // us using the same state in different workloads.
      // The sample from each iteration uses 2 random numbers.
      gen_copy.Skip((start_output * 2 * kMaxIterations +
                     Uniform::kResultElementCount - 1) /
                    Uniform::kResultElementCount);

      // Vectorized intermediate calculations for uniform rejection sampling.
      // We always generate at most 4 samples.
      Eigen::array<T, Uniform::kResultElementCount> z;
      Eigen::array<T, Uniform::kResultElementCount> g;

      const bool should_bcast = bcast.IsBroadcastingRequired();
      const auto& means_batch_indices = bcast.batch_indices(0);
      const auto& stddevs_batch_indices = bcast.batch_indices(1);
      const auto& minvals_batch_indices = bcast.batch_indices(2);
      const auto& maxvals_batch_indices = bcast.batch_indices(3);
      auto output_flat = output.data();

      // We partition work across batches and then across samples
      // per batch member, to avoid extra work.
      for (int64 output_idx = start_output; output_idx < limit_output;
           // output_idx is incremented with the inner loops below.
      ) {
        int64 batch_idx = output_idx / samples_per_batch;
        // The output layout is [samples_per_batch, num_batches]. Thus
        // the output address is sample_idx * num_batches + batch_idx.
        // Below, code will index at output_batch_offset[sample_idx *
        // num_batches] matching this.
        T* const output_batch_offset = output_flat + batch_idx;
        // Generate batch counts from BCast, as it has the right indices to loop
        // over.
        T mean, stddev, minval, maxval;
        if (should_bcast) {
          mean = means(means_batch_indices[batch_idx]);
          stddev = stddevs(stddevs_batch_indices[batch_idx]);
          minval = minvals(minvals_batch_indices[batch_idx]);
          maxval = maxvals(maxvals_batch_indices[batch_idx]);
        } else {
          mean = means(batch_idx);
          stddev = stddevs(batch_idx);
          minval = minvals(batch_idx);
          maxval = maxvals(batch_idx);
        }

        // On GPU, this check will just fill samples with NAN if it fails.
        OP_REQUIRES(ctx,
                    stddev > T(0) && minval < maxval &&
                        (Eigen::numext::isfinite(minval) ||
                         Eigen::numext::isfinite(maxval)),
                    errors::InvalidArgument("Invalid parameters"));

        int num_iterations = 0;

        // If possible, make one-sided bound be the lower bound, or make both
        // bounds positive. Otherwise, the bounds are on either side of the
        // mean.
        if ((Eigen::numext::isinf(minval) && minval < T(0)) || maxval < mean) {
          // Reverse all calculations. normMin and normMax will be flipped.
          std::swap(minval, maxval);
          stddev = -stddev;
        }

        // Calculate normalized samples, then convert them.
        const T normMin = (minval - mean) / stddev;
        const T normMax = (maxval - mean) / stddev;

        // Determine the method to use.
        const T sqrtFactor = Eigen::numext::sqrt((normMin * normMin) + T(4));
        const T cutoff =
            T(2) *
            Eigen::numext::exp(T(0.5) +
                               (normMin * (normMin - sqrtFactor)) / T(4)) /
            (normMin + sqrtFactor);
        const T diff = normMax - normMin;

        if (((normMin < -kStdDevsInsideBoundsToUseRandnSampler) &&
             (normMax >= T(0.))) ||
            ((normMax > kStdDevsInsideBoundsToUseRandnSampler) &&
             (normMin <= T(0.)))) {
          // If the bounds are a least 3 standard deviations from the mean
          // on at least one side then we rejection sample by sampling
          // from the normal distribution and rejecting samples outside
          // the bounds.
          // Under this condition the acceptance rate per iteration should
          // always be ~ 50%. This sampler is more efficient (and more
          // numerically stable when one or both bounds is far from the mean).
          for (int64 sample_idx = output_idx % samples_per_batch;
               sample_idx < samples_per_batch && output_idx < limit_output;) {
            const auto randn_sample = normal_dist(&gen_copy);
            const int size = randn_sample.size();
            for (int i = 0; i < size; ++i) {
              if ((randn_sample[i] >= normMin) &&
                  (randn_sample[i] <= normMax)) {
                output_batch_offset[sample_idx * num_batches] =
                    randn_sample[i] * stddev + mean;
                ++sample_idx;
                ++output_idx;
                if (sample_idx >= samples_per_batch ||
                    output_idx >= limit_output) {
                  break;
                }
                num_iterations = 0;
              } else {
                ++num_iterations;
                if (num_iterations > kMaxIterations) {
                  // This should never occur because this sampler should
                  // (by the selection criteria above) be used if at least 3
                  // standard deviations of one side of the distribution
                  // is within the limits (so acceptance probability per
                  // iterations >~ 1/2 per iteration).
                  LOG(ERROR) << "TruncatedNormal randn rejection sampler "
                             << "exceeded maximum iterations for "
                             << "normMin=" << normMin << " normMax=" << normMax
                             << " kMaxIterations=" << kMaxIterations;
                  ctx->SetStatus(errors::Internal(
                      "TruncatedNormal randn rejection sampler failed to accept"
                      " a sample."));
                  return;
                }
              }
            }
          }
        } else if (diff < cutoff) {
          // Sample from a uniform distribution on [normMin, normMax].

          const T plusFactor = (normMin < T(0)) ? T(0) : normMin * normMin;

          for (int64 sample_idx = output_idx % samples_per_batch;
               sample_idx < samples_per_batch && output_idx < limit_output;) {
            const auto rand = dist(&gen_copy);
            const int size = rand.size();
            // NOTE(ringwalt): These loops seem to only generate packed AVX
            // instructions for float32.
            for (int i = 0; i < size; i++) {
              z[i] = rand[i] * diff + normMin;
              g[i] = (plusFactor - z[i] * z[i]) / T(2.0);
            }

            const auto u = dist(&gen_copy);
            for (int i = 0; i < size; i++) {
              auto accept = u[i] <= Eigen::numext::exp(g[i]);
              if (accept || num_iterations + 1 >= kMaxIterations) {
                // Accept the sample z.
                // If we run out of iterations, just use the current uniform
                // sample, but emit a warning.
                // TODO(jjhunt) For small entropies (relative to the bounds),
                // this sampler is poor and may take many iterations since
                // the proposal distribution is the uniform distribution
                // U(lower_bound, upper_bound).
                if (!accept) {
                  LOG(ERROR) << "TruncatedNormal uniform rejection sampler "
                             << "exceeded max iterations. Sample may contain "
                             << "outliers.";
                  ctx->SetStatus(errors::Internal(
                      "TruncatedNormal uniform rejection sampler failed to "
                      " accept a sample."));
                  return;
                }
                output_batch_offset[sample_idx * num_batches] =
                    z[i] * stddev + mean;
                ++sample_idx;
                ++output_idx;
                if (sample_idx >= samples_per_batch ||
                    output_idx >= limit_output) {
                  break;
                }
                num_iterations = 0;
              } else {
                num_iterations++;
              }
            }
          }
        } else {
          // Sample from an exponential distribution with alpha maximizing
          // acceptance probability, offset by normMin from the origin.
          // Accept only if less than normMax.
          const T alpha =
              (normMin + Eigen::numext::sqrt((normMin * normMin) + T(4))) /
              T(2);
          for (int64 sample_idx = output_idx % samples_per_batch;
               sample_idx < samples_per_batch && output_idx < limit_output;) {
            auto rand = dist(&gen_copy);
            const int size = rand.size();
            int i = 0;
            while (i < size) {
              const T z = -Eigen::numext::log(rand[i]) / alpha + normMin;
              i++;
              const T x = normMin < alpha ? alpha - z : normMin - alpha;
              const T g = Eigen::numext::exp(-x * x / T(2.0));
              const T u = rand[i];
              i++;
              auto accept = (u <= g && z < normMax);
              if (accept || num_iterations + 1 >= kMaxIterations) {
                if (!accept) {
                  LOG(ERROR) << "TruncatedNormal exponential distribution "
                             << "rejection sampler exceeds max iterations. "
                             << "Sample may contain outliers.";
                  ctx->SetStatus(errors::Internal(
                      "TruncatedNormal exponential distribution rejection"
                      " sampler failed to accept a sample."));
                  return;
                }
                output_batch_offset[sample_idx * num_batches] =
                    z * stddev + mean;
                ++sample_idx;
                ++output_idx;
                if (sample_idx >= samples_per_batch ||
                    output_idx >= limit_output) {
                  break;
                }
                num_iterations = 0;
              } else {
                num_iterations++;
              }
            }
          }
        }
      }
    };
    // The cost of the initial calculations for the batch.
    const int64 batchInitCost =
        // normMin, normMax
        (Eigen::TensorOpCost::AddCost<T>() +
         Eigen::TensorOpCost::MulCost<T>()) *
            2
        // sqrtFactor
        + Eigen::TensorOpCost::AddCost<T>() +
        Eigen::TensorOpCost::MulCost<T>() +
        Eigen::internal::functor_traits<
            Eigen::internal::scalar_sqrt_op<T>>::Cost
        // cutoff
        + Eigen::TensorOpCost::MulCost<T>() * 4 +
        Eigen::internal::functor_traits<Eigen::internal::scalar_exp_op<T>>::Cost
        // diff
        + Eigen::TensorOpCost::AddCost<T>();
    const int64 uniformSampleCost =
        random::PhiloxRandom::kElementCost +
        random::UniformDistribution<random::PhiloxRandom, T>::kElementCost;
    // The cost of a single uniform sampling round.
    const int64 uniformRejectionSamplingCost =
        uniformSampleCost + Eigen::TensorOpCost::MulCost<T>() +
        Eigen::TensorOpCost::AddCost<T>() +
        Eigen::TensorOpCost::MulCost<T>() * 2 +
        Eigen::TensorOpCost::AddCost<T>() + uniformSampleCost +
        Eigen::internal::functor_traits<
            Eigen::internal::scalar_exp_op<T>>::Cost +
        Eigen::TensorOpCost::MulCost<T>() + Eigen::TensorOpCost::AddCost<T>();
    // Estimate the cost for an entire batch.
    // Assume we use uniform sampling, and accept the 2nd sample on average.
    const int64 batchCost = batchInitCost + uniformRejectionSamplingCost * 2;
    Shard(worker_threads.num_threads, worker_threads.workers, num_elements,
          batchCost, do_work);
  }