   ///   - values: Contains the value estimates for each step.
   ///   - finalValue: Estimated value at the final step.
   @inlinable
-  // TODO: @differentiable
   public func callAsFunction<Scalar: TensorFlowFloatingPoint>(
-    stepKinds: [Tensor<Int32>],
-    returns: [Tensor<Scalar>],
-    values: [Tensor<Scalar>],
+    stepKinds: StepKind,
+    returns: Tensor<Scalar>,
+    values: Tensor<Scalar>,
     finalValue: Tensor<Scalar>
-  ) -> [Tensor<Scalar>] {
-    precondition(stepKinds.count == returns.count && stepKinds.count == values.count)
-    return zip(returns, values).map{ $0 - $1 }
+  ) -> Tensor<Scalar> {
+    returns - values
   }
 }
 

  ///   - values: Contains the value estimates for each step.
  ///   - finalValue: Estimated value at the final step.
  @inlinable
  // TODO: @differentiable
  public func callAsFunction<Scalar: TensorFlowFloatingPoint>(
    stepKinds: [Tensor<Int32>],
    returns: [Tensor<Scalar>],
    values: [Tensor<Scalar>],
    finalValue: Tensor<Scalar>
  ) -> [Tensor<Scalar>] {
    precondition(stepKinds.count == returns.count && stepKinds.count == values.count)

    if stepKinds.isEmpty {
      return [Tensor<Scalar>]()
    }

    let T = stepKinds.count
    var advantages = [Tensor<Scalar>]()
    advantages.reserveCapacity(T)
    for t in (0..<T).reversed() {
      let futureAdvantage = t + 1 < T ?
        values[t + 1] + Scalar(discountWeight) * advantages[T - t - 2] :
        finalValue
      let discountedFutureAdvantage = Scalar(discountFactor) * futureAdvantage
      let advantage = returns[t] - values[t] + discountedFutureAdvantage.replacing(
        with: Tensor<Scalar>(zerosLike: discountedFutureAdvantage),
        where: stepKinds[t] .== StepKind.last.rawValue.scalar!)
      advantages.append(advantage)
    }
    return advantages.reversed()
  }
}

///
/// - Returns: Array of discounted return values over time.
@inlinable
public func discountedReturns<Scalar: TensorFlowNumeric>(
  discountFactor: Scalar,
  stepKinds: StepKind,
  rewards: Tensor<Scalar>,
  finalValue: Tensor<Scalar>? = nil
) -> Tensor<Scalar> {
  let isLast = stepKinds.isLast()
  let T = stepKinds.rawValue.shape[0]
  let finalReward = finalValue ?? Tensor<Scalar>(zerosLike: rewards[0])
  var discountedReturns = [Tensor<Scalar>]()
  for t in 0..<T {
    let futureReturn = T - t < T ? discountedReturns[t - 1] : finalReward
    let discountedFutureReturn = discountFactor * futureReturn
    let discountedReturn = rewards[T - t - 1] + discountedFutureReturn.replacing(
      with: Tensor<Scalar>(zerosLike: discountedFutureReturn),
      where: isLast[T - t - 1])
    discountedReturns.append(discountedReturn)
  }
  return Tensor<Scalar>(stacking: discountedReturns.reversed())

