         return dtype_dispatcher(dtype_id, _cub_segmented_reduce_max(),
                    workspace, workspace_size, x, y, num_segments, offset_start,
                    offset_end, stream);
+    case CUPY_CUB_PROD:
+        return dtype_dispatcher(dtype_id, _cub_segmented_reduce_prod(),
+                   workspace, workspace_size, x, y, num_segments, offset_start,
+                   offset_end, stream);
     default:
         throw std::runtime_error("Unsupported operation");
     }

    }
};

//struct _cub_segmented_reduce_prod {
//    template <typename T>
//    void operator()(void* workspace, size_t& workspace_size, void* x, void* y,
//        int num_segments, void* offset_start, void* offset_end, cudaStream_t s)
//    {
//        DeviceSegmentedReduce::Sum(workspace, workspace_size,
//            static_cast<T*>(x), static_cast<T*>(y), num_segments,
//            static_cast<int*>(offset_start),
//            static_cast<int*>(offset_end), s);
//    }
//};

//
// **** CUB Min ****



/* ------------------------------------ "Patches" to CUB ------------------------------------
   This stub is needed because CUB does not have a built-in "prod" operator
*/

//
// product functor
//
struct _multiply
{
    template <typename T>
    __host__ __device__ __forceinline__ T operator()(const T &a, const T &b) const
    {
        return a * b;
    }
};

/*
   These stubs are needed because CUB does not handle NaNs properly, while NumPy has certain
   behaviors with which we must comply.
   TODO(leofang): support half precision?
*/

//
// Max()
