Fix _to_values_offsets handle converting 1d values

This function is now used in more cases than it used to be and needs to be generalized a bit.

Author: karlhigley

merlin/systems/triton/conversions.py

@@ -111,12 +111,19 @@ def _to_values_offsets(array):
     values, offsets
         Tuple of values and offsets
     """
-    num_rows = array.shape[0]
-    row_lengths = [array.shape[1]] * num_rows
-    offsets = [0] + list(itertools.accumulate(row_lengths))
     array_lib = cp if cp and isinstance(array, cp.ndarray) else np
+    shape = array.shape
+
+    values_shape = [-1]
+    if len(shape) > 2:
+        values_shape.extend(shape[2:])
+    values = array.reshape(*values_shape)
+
+    num_rows = shape[0]
+    row_lengths = [shape[1]] * num_rows if len(shape) > 1 else [1]
+    offsets = [0] + list(itertools.accumulate(row_lengths))
     offsets = array_lib.array(offsets, dtype="int32")
-    values = array.reshape(-1, *array.shape[2:])
+
     return values, offsets