[plotting] introduces work aroung for plotting of single variable tensors

plotting.py

@@ -280,7 +280,7 @@ def figure_inputs(
     rows, cols = multiplot.lenghts()
     size = 5
     fig, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
-
+    ax = ax if isinstance(ax, np.ndarray) else np.array([ax])
     inps = inps.reshape(inps.shape[0], -1)
     inps = np.clip(inps, -overflow, overflow)
     order = np.argsort(-(sample_weight[:, None] * truth).sum(axis=0))
@@ -337,6 +337,7 @@ def figure_weight_study(
     rows, cols = multiplot.lenghts()
     size = 5
     fig, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
+    ax = ax if isinstance(ax, np.ndarray) else np.array([ax])
 
     class_inps = [inps.reshape(inps.shape[0], -1) for inps in class_inps]
     for feat, name in enumerate(columns):