BUG: 0.45.0 update breaks pytorch example on docs

[hongfanmeng]

Issue Description

Update 0.45.0 has a breaking change

• Changed type and shape of returned SHAP values ​​in some cases, to be consistent with model outputs. SHAP values ​​for models with multiple outputs are now np.ndarray rather than list, by @CloseChoice in BREAKING: make shap explanations consistent for xgboost and lightgbm #3318

After the update, the pytorch example in the documentation process shap_values ​​incorrectly
https://shap.readthedocs.io/en/latest/example_notebooks/image_examples/image_classification/PyTorch%20Deep%20Explainer%20MNIST%20example.html

shap_numpy = [np.swapaxes(np.swapaxes(s, 1, -1), 1, 2) for s in shap_values]
test_numpy = np.swapaxes(np.swapaxes(test_images.numpy(), 1, -1), 1, 2)

Causing an error in the following code

# plot the feature attributions
shap.image_plot(list(shap_numpy), -test_numpy)

Updating the processing of output to the following version should resolve the issue

# shap_numpy = [np.swapaxes(np.swapaxes(s, 1, -1), 1, 2) for s in shap_values]
shap_numpy = list(np.transpose(shap_values, (4, 0, 2, 3, 1)))
test_numpy = np.swapaxes(np.swapaxes(test_images.numpy(), 1, -1), 1, 2)

Minimal Reproducible Example

See above

Traceback

IndexError                                Traceback (most recent call last)
Cell In[17], line 2
      1 # plot the feature attributions
----> 2 shap.image_plot(list(shap_numpy), -test_numpy)

File ~/miniconda3/envs/iiot/lib/python3.11/site-packages/shap/plots/_image.py:154, in image(shap_values, pixel_values, labels, true_labels, width, aspect, hspace, labelpad, cmap, show)
    152     axes[row, 0].set_title(true_labels[row], **label_kwargs)
    153 axes[row, 0].axis('off')
--> 154 if len(shap_values[0][row].shape) == 2:
    155     abs_vals = np.stack([np.abs(shap_values[i]) for i in range(len(shap_values))], 0).flatten()
    156 else:

IndexError: index 1 is out of bounds for axis 0 with size 1

Expected Behavior

No response

Bug report checklist

• I have checked that this issue has not already been reported.
• I have confirmed this bug exists on the latest release of shap.
• I have confirmed this bug exists on the master branch of shap.
• I'd be interested in making a PR to fix this bug

Installed Versions

0.45.0

[CloseChoice]

Thanks for reporting and a special thanks for figuring out a solution. I created a PR using exactly the code you suggested.