WebGPU backend fails to run DeepLab V3 model

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No
• TensorFlow.js installed from (npm or script link): build tfjs-webgpu-backend from source
• TensorFlow.js version (use command below): 3.14.0
• Browser version: Chrome 99.0.4844.74
• Tensorflow.js Converter Version: use models in stock

Describe the current behavior Build tfjs-webgpu-backend from source and import it. Download tfjs deeplab v3 model from: https://tfhub.dev/tensorflow/tfjs-model/deeplab/pascal/1/default/1

await tf.setBackend('webgpu');
const deeplab = await tf.loadGraphModel('./deeplab_pascal_1_default_1/model.json');
const input = tf.zeros([1, 513, 513, 3], 'int32');
deeplab.predict(input);

It throws error in console as

Uncaught Error: Unsupported 6D shape
    at Object.Uv [as assert] (util_base.js:153:11)
    at backend_webgpu.ts:648:16
    at Array.forEach (<anonymous>)
    at WebGPUBackend.makeUniforms (backend_webgpu.ts:628:22)
    at WebGPUBackend.runWebGPUProgram (backend_webgpu.ts:750:27)
    at transpose (Transpose.ts:52:24)
    at Object.spaceToBatchND [as kernelFunc] (SpaceToBatchND.ts:68:20)
    at n (engine.js:644:22)
    at engine.js:711:23
    at e.t.scopedRun (engine.js:478:19)

The webgl backend works fine for this model.

Describe the expected behavior

TF.js webgpu backend is able to run this model.

/cc @qjia7 @pyu10055 @lina128

[qjia7]

@haoyunfeix is working on this. #6321

[huningxin]

Sorry for late update. This issue is fixed according to my test. Thanks again for the resolution.