   ///   - 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
   }
 }
 

  ///   - 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.

  public init() {}

  @inlinable
  public func callAsFunction<Scalar: TensorFlowFloatingPoint>(
    stepKinds: StepKind,
    returns: Tensor<Scalar>,
    values: Tensor<Scalar>,
    finalValue: Tensor<Scalar>
  ) -> Tensor<Scalar> {
    returns
  }
}
