   public init() {}
 
   @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>] {
+  ) -> Tensor<Scalar> {
     returns
   }
 }

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

// TODO: Remove once optionals are differentiable.

///
/// - 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())

