DispatcherReduceByKey uses the scheduling_range like all other dispatchers

vtkm/worklet/DispatcherReduceByKey.h

@@ -59,8 +59,8 @@ public:
     const InputDomainType& inputDomain = invocation.GetInputDomain();
 
     // Now that we have the input domain, we can extract the range of the
-    // scheduling and call BadicInvoke.
-    this->BasicInvoke(invocation, inputDomain.GetInputRange());
+    // scheduling and call BasicInvoke.
+    this->BasicInvoke(invocation, internal::scheduling_range(inputDomain));
   }
 };
 }

vtkm/worklet/internal/DispatcherBase.h

@@ -48,6 +48,8 @@ namespace vtkm
 {
 namespace worklet
 {
+template <typename T>
+class VTKM_ALWAYS_EXPORT Keys;
 namespace internal
 {
 
@@ -57,6 +59,13 @@ inline auto scheduling_range(const Domain& inputDomain) -> decltype(inputDomain.
   return inputDomain.GetNumberOfValues();
 }
 
+template <typename KeyType>
+inline auto scheduling_range(const vtkm::worklet::Keys<KeyType>& inputDomain)
+  -> decltype(inputDomain.GetInputRange())
+{
+  return inputDomain.GetInputRange();
+}
+
 template <typename Domain>
 inline auto scheduling_range(const Domain* const inputDomain)
   -> decltype(inputDomain->GetNumberOfValues())
@@ -64,6 +73,13 @@ inline auto scheduling_range(const Domain* const inputDomain)
   return inputDomain->GetNumberOfValues();
 }
 
+template <typename KeyType>
+inline auto scheduling_range(const vtkm::worklet::Keys<KeyType>* const inputDomain)
+  -> decltype(inputDomain->GetInputRange())
+{
+  return inputDomain->GetInputRange();
+}
+
 template <typename Domain, typename SchedulingRangeType>
 inline auto scheduling_range(const Domain& inputDomain, SchedulingRangeType type)
   -> decltype(inputDomain.GetSchedulingRange(type))