Cannot run a existing model training in Tensorflow and converted by converter, `transpose` rank error

[Lutra-Fs]

Please make sure that this is a bug. As per our
GitHub Policy,
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 11 23590.1000 and WSL Arch
• TensorFlow.js installed from (npm or script link): npm
• TensorFlow.js version (use command below): 4.13.0
• Browser version: Microsoft Edge 121.0.2220.3
• Tensorflow.js Converter Version: 4.13.0
• test node js version: 18.17.1

Describe the current behavior

Error: Error in transpose: rank of input 2 must match length of perm 0,2,3,1.
    at assert (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:7748:15)
    at transpose_ (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:23465:5)
    at Object.transpose__op [as transpose] (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:12220:29)
    at executeOp$9 (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:29599:25)
    at F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:30175:50
    at F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-core\dist\tf-core.node.js:4562:22
    at Engine.scopedRun (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-core\dist\tf-core.node.js:4572:23)
    at Engine.tidy (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-core\dist\tf-core.node.js:4561:21)
    at tidy (F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-core\dist\tf-core.node.js:8307:19)
    at F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:30175:24

when running rfft and the after transpose in the browser using WebGL:
in

if ($x.dtype === 'complex64') {
    return tidy(() => {
      let $real = real($x);
      let $imag = imag($x);
      $real = ENGINE.runKernel(
          Transpose, {x: $real} as unknown as NamedTensorMap,
          attrs as unknown as NamedAttrMap);
      $imag = ENGINE.runKernel(
          Transpose, {x: $imag} as unknown as NamedTensorMap,
          attrs as unknown as NamedAttrMap);
      if (conjugate) {
        $imag = neg($imag);
      }
      return complex($real, $imag);
    });
  } 

in transpose.ts in @tfjs/core, although x is a complex64 shape in [1, 13, 64, 33], in let $real = real($x); let $imag = imag($x);
these two functions get a tensor with shape [832,33]. I have tried to find where the issue happens in the code but failed.

Describe the expected behavior
model runs correctly
Standalone code to reproduce the issue

const tf = require('@tensorflow/tfjs');

tf.setBackend('cpu');
tf.enableDebugMode();
tf.tidy(async () => {
  const tfmodel = await tf.loadGraphModel(
    'http://localhost:3000/FNO2dLite_from_saved_model/model.json',
  );

  const inputTensor = tf.zeros([1, 64, 64, 5]);

  const output = await tfmodel.executeAsync(inputTensor);
});

converted model:
FNO2dLite_from_saved_model.zip

Other info / logs

Node.js cpu backend debug info:

Cast                            0.2402999997138977ms    1D 3                    3       x: 1D 3 
GatherV2                        0.8700000047683716ms    1D 3                    3       x: 1D 4 indices: 1D 3
Concat                          0.3334999978542328ms    1D 4                    4       0: 1D 3 1: 1D 1
Prod                            0.2727999985218048ms    0D                      1       x: 1D 3
Cast                            0.017899997532367706ms  1D 1                    1       x: 1D 1
GatherV2                        0.05390000343322754ms   1D 1                    1       x: 1D 4 indices: 1D 1
Prod                            0.024800002574920654ms  0D                      1       x: 1D 1
Reshape                         0.022599995136260986ms  0D                      1       x: 0D
Pack                            0.3401999995112419ms    1D 2                    2       0: 0D  1: 0D
Reshape                         0.01990000158548355ms   2D 4096,5               20480   x: 4D 1,64,64,5
BatchMatMul                     9.050800003111362ms     2D 4096,13              53248   a: 2D 4096,5 b: 2D 5,13
Reshape                         0.014400005340576172ms  4D 1,64,64,13           53248   x: 2D 4096,13
Add                             3.8828999996185303ms    4D 1,64,64,13           53248   a: 4D 1,64,64,13 b: 1D 13
StridedSlice                    0.8957000002264977ms    0D                      1       x: 1D 4
Reshape                         0.0076000019907951355ms 0D                      1       x: 0D
Pack                            0.2149999961256981ms    1D 5                    5       0: 0D  1: 0D  2: 0D  3: 0D  4: 0D
Fill                            0.15600000321865082ms   5D 1,13,64,25,2         41600
StridedSlice                    0.12000000476837158ms   0D                      1       x: 1D 4
Reshape                         0.00849999487400055ms   0D                      1       x: 0D
Pack                            0.1729000061750412ms    1D 5                    5       0: 0D  1: 0D  2: 0D  3: 0D  4: 0D
Fill                            0.03220000118017197ms   5D 1,13,48,8,2          9984
Transpose                       9.9307000041008ms       4D 1,13,64,64           53248   x: 4D 1,64,64,13
ZerosLike                       0.17249999940395355ms   4D 1,13,64,64           53248   x: 4D 1,13,64,64
Complex                         0.0738999992609024ms    4D 1,13,64,64           53248   real: 4D 1,13,64,64 imag: 4D 1,13,64,64
Reshape                         0.01759999990463257ms   2D 832,64               53248   x: 4D 1,13,64,64
FFT                             2488.2494000047445ms    2D 832,64               53248   input: 2D 832,64
Real                            0.10300000011920929ms   2D 832,64               53248   input: 2D 832,64
Imag                            0.08919999748468399ms   2D 832,64               53248   input: 2D 832,64
SplitV                          67.00439999997616ms     2D 832,33               27456   x: 2D 832,64
SplitV                          67.00439999997616ms     2D 832,31               25792   x: 2D 832,64
SplitV                          44.32190000265837ms     2D 832,33               27456   x: 2D 832,64
SplitV                          44.32190000265837ms     2D 832,31               25792   x: 2D 832,64
Complex                         0.12290000170469284ms   2D 832,33               27456   real: 2D 832,33 imag: 2D 832,33
Reshape                         0.07990000396966934ms   4D 1,13,64,33           27456   x: 2D 832,33
Real                            0.014399997889995575ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Imag                            0.006999999284744263ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Transpose                       1.7778000012040138ms    4D 1,13,33,64           27456   x: 4D 1,13,64,33
Transpose                       9.873700000345707ms     4D 1,13,33,64           27456   x: 4D 1,13,64,33
Complex                         0.012400001287460327ms  4D 1,13,33,64           27456   real: 4D 1,13,33,64 imag: 4D 1,13,33,64
FFT                             1274.7540999948978ms    4D 1,13,33,64           27456   input: 4D 1,13,33,64
Real                            0.008600004017353058ms  4D 1,13,33,64           27456   input: 4D 1,13,33,64
Imag                            0.0054000020027160645ms 4D 1,13,33,64           27456   input: 4D 1,13,33,64
Transpose                       1.6305999979376793ms    4D 1,13,64,33           27456   x: 4D 1,13,33,64
Transpose                       6.2162000015378ms       4D 1,13,64,33           27456   x: 4D 1,13,33,64
Complex                         0.011500000953674316ms  4D 1,13,64,33           27456   real: 4D 1,13,64,33 imag: 4D 1,13,64,33
Imag                            0.009499996900558472ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Real                            0.007199995219707489ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Reshape                         0.004900000989437103ms  3D 13,64,33             27456   x: 4D 1,13,64,33
Pack                            4.126200005412102ms     5D 1,13,64,33,2         54912   0: 4D 1,13,64,33 1: 4D 1,13,64,33
StridedSlice                    2.8289000019431114ms    5D 1,13,8,8,2           1664    x: 5D 1,13,64,33,2
StridedSlice                    2.217600002884865ms     4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          23.340500004589558ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
StridedSlice                    1.3998000025749207ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          4.960900001227856ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
Add                             0.10509999841451645ms   4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8
StridedSlice                    1.5082999989390373ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          4.740699999034405ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
StridedSlice                    1.7577999979257584ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          3.125200003385544ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
Sub                             0.085999995470047ms     4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8
Reshape                         0.004700005054473877ms  3D 13,8,8               832     x: 4D 1,13,8,8
Pack                            1.6943999975919724ms    5D 1,13,8,8,2           1664    0: 4D 1,13,8,8 1: 4D 1,13,8,8
StridedSlice                    1.412200003862381ms     5D 1,13,8,8,2           1664    x: 5D 1,13,64,33,2
StridedSlice                    0.9057000055909157ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          2.745699994266033ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
StridedSlice                    0.6583999991416931ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          2.7393999993801117ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
Add                             0.13439999520778656ms   4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8
StridedSlice                    2.185900002717972ms     4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          2.4505999982357025ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
StridedSlice                    0.6269000023603439ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2
Einsum                          2.419700004160404ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8
Sub                             0.5145999938249588ms    4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8
Reshape                         0.005199998617172241ms  3D 13,8,8               832     x: 4D 1,13,8,8
Pack                            0.042100004851818085ms  5D 1,13,8,8,2           1664    0: 4D 1,13,8,8 1: 4D 1,13,8,8
Concat                          0.15489999949932098ms   5D 1,13,64,8,2          13312   0: 5D 1,13,8,8,2 1: 5D 1,13,48,8,2 2: 5D 1,13,8,8,2
Concat                          0.13910000026226044ms   5D 1,13,64,33,2         54912   0: 5D 1,13,64,8,2 1: 5D 1,13,64,25,2
StridedSlice                    19.49159999936819ms     4D 1,13,64,33           27456   x: 5D 1,13,64,33,2
StridedSlice                    20.73200000077486ms     4D 1,13,64,33           27456   x: 5D 1,13,64,33,2
Complex                         0.010799996554851532ms  4D 1,13,64,33           27456   real: 4D 1,13,64,33 imag: 4D 1,13,64,33
Real                            0.008000001311302185ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Imag                            0.006200000643730164ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Transpose                       6.112900003790855ms     4D 1,13,33,64           27456   x: 4D 1,13,64,33
Transpose                       9.36089999973774ms      4D 1,13,33,64           27456   x: 4D 1,13,64,33
Complex                         0.010499998927116394ms  4D 1,13,33,64           27456   real: 4D 1,13,33,64 imag: 4D 1,13,33,64
IFFT                            1203.483599998057ms     4D 1,13,33,64           27456   input: 4D 1,13,33,64
Real                            0.007899999618530273ms  4D 1,13,33,64           27456   input: 4D 1,13,33,64
Imag                            0.00469999760389328ms   4D 1,13,33,64           27456   input: 4D 1,13,33,64
Transpose                       1.4570000022649765ms    4D 1,13,64,33           27456   x: 4D 1,13,33,64
Transpose                       1.4847000017762184ms    4D 1,13,64,33           27456   x: 4D 1,13,33,64
Complex                         0.006900005042552948ms  4D 1,13,64,33           27456   real: 4D 1,13,64,33 imag: 4D 1,13,64,33
Real                            0.006599999964237213ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33
Reshape                         0.0040000006556510925ms 2D 832,33               27456   x: 4D 1,13,64,33
Imag                            0.00509999692440033ms   4D 1,13,64,33           27456   input: 4D 1,13,64,33
Reshape                         0.002499997615814209ms  2D 832,33               27456   x: 4D 1,13,64,33
Slice                           25.670800000429153ms    2D 832,31               25792   x: 2D 832,33
Reverse                         18.266600005328655ms    2D 832,31               25792   x: 2D 832,31
Slice                           25.87559999525547ms     2D 832,31               25792   x: 2D 832,33
Reverse                         27.317199997603893ms    2D 832,31               25792   x: 2D 832,31
Multiply                        1.4156000018119812ms    2D 832,31               25792   a: 2D 832,31 b: 0D
Concat                          0.4617999941110611ms    2D 832,64               53248   0: 2D 832,33 1: 2D 832,31
Concat                          0.16420000046491623ms   2D 832,64               53248   0: 2D 832,33 1: 2D 832,31
Complex                         0.006299994885921478ms  2D 832,64               53248   real: 2D 832,64 imag: 2D 832,64
Reshape                         0.0054000020027160645ms 2D 832,64               53248   x: 2D 832,64
IFFT                            2520.780900001526ms     2D 832,64               53248   input: 2D 832,64
Real                            0.009300000965595245ms  2D 832,64               53248   input: 2D 832,64
F:\GitRepos\CFlowSim\node_modules\@tensorflow\tfjs\node_modules\@tensorflow\tfjs-converter\dist\tf-converter.node.js:7748
        throw new Error(typeof msg === 'string' ? msg : msg());
              ^

Error: Error in transpose: rank of input 2 must match length of perm 0,2,3,1.

[gaikwadrahul8]

Hi, @Lutra-Fs

I apologize for the delayed response and I've invested time in replicating the reported issue on my end and have observed the same behavior you've described. It's evident that further investigation is necessary to thoroughly understand and address this issue. Thank you for bringing this issue to our attention and will update you soon.

Thank you for your understanding and patience.

Here is error output log for reference :

(base) gaikwadrahul-macbookpro:test-8084 gaikwadrahul$ node index.js 
Debugging mode is ON. The output of every math call will be downloaded to CPU and checked for NaNs. This significantly impacts performance.

============================
Hi, looks like you are running TensorFlow.js in Node.js. To speed things up dramatically, install our node backend, visit https://github.com/tensorflow/tfjs-node for more details. 
============================
Cast                            0.1474170684814453ms    1D 3                    3       x: 1D 3 
GatherV2                        0.4168750047683716ms    1D 3                    3       x: 1D 4 indices: 1D 3 
Concat                          0.17524993419647217ms   1D 4                    4       0: 1D 3 1: 1D 1 
Prod                            0.15112507343292236ms   0D                      1       x: 1D 3 
Cast                            0.010084033012390137ms  1D 1                    1       x: 1D 1 
GatherV2                        0.037457942962646484ms  1D 1                    1       x: 1D 4 indices: 1D 1 
Prod                            0.02545905113220215ms   0D                      1       x: 1D 1 
Reshape                         0.008957982063293457ms  0D                      1       x: 0D  
Pack                            0.21616697311401367ms   1D 2                    2       0: 0D  1: 0D  
Reshape                         0.010916948318481445ms  2D 4096,5               20480   x: 4D 1,64,64,5 
BatchMatMul                     5.396500110626221ms     2D 4096,13              53248   a: 2D 4096,5 b: 2D 5,13 
Reshape                         0.007834076881408691ms  4D 1,64,64,13           53248   x: 2D 4096,13 
Add                             2.291707992553711ms     4D 1,64,64,13           53248   a: 4D 1,64,64,13 b: 1D 13 
StridedSlice                    0.48170900344848633ms   0D                      1       x: 1D 4 
Reshape                         0.003416895866394043ms  0D                      1       x: 0D  
Pack                            0.17920804023742676ms   1D 5                    5       0: 0D  1: 0D  2: 0D  3: 0D  4: 0D  
Fill                            0.05324995517730713ms   5D 1,13,64,25,2         41600
StridedSlice                    0.11483395099639893ms   0D                      1       x: 1D 4 
Reshape                         0.0021671056747436523ms 0D                      1       x: 0D  
Pack                            0.05204200744628906ms   1D 5                    5       0: 0D  1: 0D  2: 0D  3: 0D  4: 0D  
Fill                            0.0071250200271606445ms 5D 1,13,48,8,2          9984
Transpose                       7.802291989326477ms     4D 1,13,64,64           53248   x: 4D 1,64,64,13 
ZerosLike                       0.08791697025299072ms   4D 1,13,64,64           53248   x: 4D 1,13,64,64 
Complex                         0.0307919979095459ms    4D 1,13,64,64           53248   real: 4D 1,13,64,64 imag: 4D 1,13,64,64 
Reshape                         0.01258397102355957ms   2D 832,64               53248   x: 4D 1,13,64,64 
FFT                             1082.5460410118103ms    2D 832,64               53248   input: 2D 832,64 
Real                            0.008332967758178711ms  2D 832,64               53248   input: 2D 832,64 
Imag                            0.004791975021362305ms  2D 832,64               53248   input: 2D 832,64 
SplitV                          36.385082960128784ms    2D 832,33               27456   x: 2D 832,64 
SplitV                          36.385082960128784ms    2D 832,31               25792   x: 2D 832,64 
SplitV                          17.55787491798401ms     2D 832,33               27456   x: 2D 832,64 
SplitV                          17.55787491798401ms     2D 832,31               25792   x: 2D 832,64 
Complex                         0.0737919807434082ms    2D 832,33               27456   real: 2D 832,33 imag: 2D 832,33 
Reshape                         0.030874967575073242ms  4D 1,13,64,33           27456   x: 2D 832,33 
Real                            0.003916025161743164ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Imag                            0.0028330087661743164ms 4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Transpose                       2.2252920866012573ms    4D 1,13,33,64           27456   x: 4D 1,13,64,33 
Transpose                       7.788125038146973ms     4D 1,13,33,64           27456   x: 4D 1,13,64,33 
Complex                         0.007250070571899414ms  4D 1,13,33,64           27456   real: 4D 1,13,33,64 imag: 4D 1,13,33,64 
FFT                             543.9830000400543ms     4D 1,13,33,64           27456   input: 4D 1,13,33,64 
Real                            0.0048749446868896484ms 4D 1,13,33,64           27456   input: 4D 1,13,33,64 
Imag                            0.0034159421920776367ms 4D 1,13,33,64           27456   input: 4D 1,13,33,64 
Transpose                       1.3268749713897705ms    4D 1,13,64,33           27456   x: 4D 1,13,33,64 
Transpose                       1.5280840396881104ms    4D 1,13,64,33           27456   x: 4D 1,13,33,64 
Complex                         0.0055419206619262695ms 4D 1,13,64,33           27456   real: 4D 1,13,64,33 imag: 4D 1,13,64,33 
Imag                            0.004707932472229004ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Real                            0.0030829906463623047ms 4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Reshape                         0.0035829544067382812ms 3D 13,64,33             27456   x: 4D 1,13,64,33 
Pack                            2.7575000524520874ms    5D 1,13,64,33,2         54912   0: 4D 1,13,64,33 1: 4D 1,13,64,33 
StridedSlice                    1.5972920656204224ms    5D 1,13,8,8,2           1664    x: 5D 1,13,64,33,2 
StridedSlice                    0.8508750200271606ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          10.04020893573761ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
StridedSlice                    0.9662920236587524ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          3.3536670207977295ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
Add                             0.060249924659729004ms  4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8 
StridedSlice                    0.40299999713897705ms   4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          2.6323330402374268ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
StridedSlice                    0.5367909669876099ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          2.2162081003189087ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
Sub                             0.06120800971984863ms   4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8 
Reshape                         0.002290964126586914ms  3D 13,8,8               832     x: 4D 1,13,8,8 
Pack                            0.9710409641265869ms    5D 1,13,8,8,2           1664    0: 4D 1,13,8,8 1: 4D 1,13,8,8 
StridedSlice                    0.8078329563140869ms    5D 1,13,8,8,2           1664    x: 5D 1,13,64,33,2 
StridedSlice                    0.6987079381942749ms    4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          2.342542052268982ms     4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
StridedSlice                    0.39966702461242676ms   4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          2.3813750743865967ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
Add                             0.05741703510284424ms   4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8 
StridedSlice                    0.784000039100647ms     4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          1.9805830717086792ms    4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
StridedSlice                    0.402042031288147ms     4D 1,13,8,8             832     x: 5D 1,13,8,8,2 
Einsum                          1.98170804977417ms      4D 1,13,8,8             832     0: 4D 1,13,8,8 1: 4D 13,13,8,8 
Sub                             0.11024999618530273ms   4D 1,13,8,8             832     a: 4D 1,13,8,8 b: 4D 1,13,8,8 
Reshape                         0.002957940101623535ms  3D 13,8,8               832     x: 4D 1,13,8,8 
Pack                            0.025207996368408203ms  5D 1,13,8,8,2           1664    0: 4D 1,13,8,8 1: 4D 1,13,8,8 
Concat                          0.08537495136260986ms   5D 1,13,64,8,2          13312   0: 5D 1,13,8,8,2 1: 5D 1,13,48,8,2 2: 5D 1,13,8,8,2 
Concat                          0.08041703701019287ms   5D 1,13,64,33,2         54912   0: 5D 1,13,64,8,2 1: 5D 1,13,64,25,2 
StridedSlice                    13.768375039100647ms    4D 1,13,64,33           27456   x: 5D 1,13,64,33,2 
StridedSlice                    13.86758303642273ms     4D 1,13,64,33           27456   x: 5D 1,13,64,33,2 
Complex                         0.007917046546936035ms  4D 1,13,64,33           27456   real: 4D 1,13,64,33 imag: 4D 1,13,64,33 
Real                            0.004541993141174316ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Imag                            0.003041982650756836ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Transpose                       4.4297919273376465ms    4D 1,13,33,64           27456   x: 4D 1,13,64,33 
Transpose                       6.424625039100647ms     4D 1,13,33,64           27456   x: 4D 1,13,64,33 
Complex                         0.008000016212463379ms  4D 1,13,33,64           27456   real: 4D 1,13,33,64 imag: 4D 1,13,33,64 
IFFT                            545.8493330478668ms     4D 1,13,33,64           27456   input: 4D 1,13,33,64 
Real                            0.005500078201293945ms  4D 1,13,33,64           27456   input: 4D 1,13,33,64 
Imag                            0.005625009536743164ms  4D 1,13,33,64           27456   input: 4D 1,13,33,64 
Transpose                       1.300374984741211ms     4D 1,13,64,33           27456   x: 4D 1,13,33,64 
Transpose                       1.5705831050872803ms    4D 1,13,64,33           27456   x: 4D 1,13,33,64 
Complex                         0.0035001039505004883ms 4D 1,13,64,33           27456   real: 4D 1,13,64,33 imag: 4D 1,13,64,33 
Real                            0.0034589767456054688ms 4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Reshape                         0.0024160146713256836ms 2D 832,33               27456   x: 4D 1,13,64,33 
Imag                            0.002624988555908203ms  4D 1,13,64,33           27456   input: 4D 1,13,64,33 
Reshape                         0.0015000104904174805ms 2D 832,33               27456   x: 4D 1,13,64,33 
Slice                           14.373375058174133ms    2D 832,31               25792   x: 2D 832,33 
Reverse                         11.420541048049927ms    2D 832,31               25792   x: 2D 832,31 
Slice                           15.864165902137756ms    2D 832,31               25792   x: 2D 832,33 
Reverse                         15.844208002090454ms    2D 832,31               25792   x: 2D 832,31 
Multiply                        1.0933750867843628ms    2D 832,31               25792   a: 2D 832,31 b: 0D  
Concat                          0.16737496852874756ms   2D 832,64               53248   0: 2D 832,33 1: 2D 832,31 
Concat                          0.10745894908905029ms   2D 832,64               53248   0: 2D 832,33 1: 2D 832,31 
Complex                         0.005750060081481934ms  2D 832,64               53248   real: 2D 832,64 imag: 2D 832,64 
Reshape                         0.002084016799926758ms  2D 832,64               53248   x: 2D 832,64 
IFFT                            1169.826917052269ms     2D 832,64               53248   input: 2D 832,64 
Real                            0.004917025566101074ms  2D 832,64               53248   input: 2D 832,64 
Error: Error: Error in transpose: rank of input 2 must match length of perm 0,2,3,1.
    at assert (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-converter/dist/tf-converter.node.js:7745:15)
    at transpose_ (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-converter/dist/tf-converter.node.js:23462:5)
    at Object.transpose__op [as transpose] (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-converter/dist/tf-converter.node.js:12217:29)
    at executeOp$9 (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-converter/dist/tf-converter.node.js:29596:25)
    at /Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-converter/dist/tf-converter.node.js:30172:50
    at /Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:4562:22
    at Engine.scopedRun (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:4572:23)
    at Engine.tidy (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:4561:21)
    at tidy (/Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-core/dist/tf-core.node.js:8307:19)
    at /Users/gaikwadrahul/Desktop/TFJS/test-8084/node_modules/@tensorflow/tfjs-node/node_modules/@tensorflow/tfjs-converter/dist/tf-converter.node.js:30172:24
(base) gaikwadrahul-macbookpro:test-8084 gaikwadrahul$ 

[Gabfm1984]

Hi, I want to post some workarounds a found while getting my audio model to work on tensorflow js. Here's the code for a keras model that converts 1D data to 2D using STFT [height, width, channels]. However batch and channels must be 1 for this to work. Also, this concatenates imag and real into the last axis.

def modelTestFTv4(fftSize=2048):
        
        #Para probar como hacer las transformaciones y si son posibles en tensorflow JS.
        #This model applies an STFT and inverse STFT on the time-domain signal.
        #v3: Concatena real-imag
        #v4: Ya prepara los datos para procesarlos
        
        options_Normalize = True
    
        # ---- Time domain Input: shape [ReverbModels.INPUTSIZE, ReverbModels.INPUTCHANNELS]
        inputs = tf.keras.layers.Input(shape=[ReverbModels.INPUTSIZE, ReverbModels.INPUTCHANNELS])
        inputFinal = inputs  # original time-domain signal
    
        # ---- Normalize
        if options_Normalize:
            target_mean = 0.3
            input_mean = tf.reduce_mean(tf.abs(inputFinal)) + 1e-7
            inputFinal = inputFinal * (target_mean / input_mean)
    
        # ---- STFT/Inverse-STFT parameters
        frame_length = fftSize
        frame_step = fftSize // 4  # 75% overlap
        window_fn = tf.signal.hann_window
        output_length = ReverbModels.INPUTSIZE
    
        x = tf.transpose(inputFinal, perm=[0,2,1])        
        # ---- Apply STFT: output shape will be [batch, channels, frames, fft_bins]
        stft_result = tf.signal.stft(
            x,
            frame_length=frame_length,
            frame_step=frame_step,
            fft_length=fftSize,
            window_fn=window_fn,
            pad_end=True
        )
        
        # ---- Get real/imag components
        _real = tf.math.real(stft_result)
        _imag = tf.math.imag(stft_result)
        
        wshape = tf.shape(_real)
        
        magSpec = tf.abs(stft_result)
        
        specShape = tf.shape(_real)
        
        # # ---- Need to reshape so it works in JS
        _real = tf.reshape(_real, [specShape[-2], specShape[-1]])
        _imag = tf.reshape(_imag, [specShape[-2], specShape[-1]])
        #_mag = tf.reshape(magSpec, [specShape[-2], specShape[-1]])

        # # ---- Concat and transpose to our HWC (remove batch dimension)
        inputData = tf.stack([_real, _imag], axis=2)
        
        
        inputData = tf.transpose(inputData, perm=[1,0,2]) # Transpose for H-W-C
                
        
        # ---- PROCESSING PART
        inputData *= 0.5
        # ---- END OF PROCESSING
        
        # ---- Split to real and imag
        inputData = tf.transpose(inputData, perm=[1,0,2]) # Transpose for W-H-C
        _real, _imag = tf.split(inputData, 2, axis=-1)
        
        # # Rebuild complex spectrum, removing the first dimension (channel), so istft works (JS)
        
                
        processedSpec = tf.complex(_real, _imag)                
        processedSpec = tf.reshape(processedSpec, [wshape[-2], wshape[-1]])
        
        # ---- Inverse STFT: reconstruct time-domain signal with specified output length
        stft_result =  processedSpec 
        
        reconstructed = tf.signal.inverse_stft(
            stft_result,
            frame_length=frame_length,
            frame_step=frame_step,
            window_fn=window_fn,
        )
        # reconstructed shape: [samples]
    
        # ---- Reshape back to original format
        # Remove extra samples
        recon_reshaped = reconstructed[..., :ReverbModels.INPUTSIZE]

        waveformX = tf.expand_dims(recon_reshaped, axis=1) # Add channel dimension
    
        # ---- DeNormalize
        if options_Normalize:
            waveformX = waveformX * (input_mean / target_mean)
    
        model_name = "modelTestFT"
        waveformX = tf.expand_dims(waveformX, axis=0) # Add fake batch dimension
        
        return tf.keras.Model(inputs=inputs, outputs=waveformX, name=model_name)

Hope this helps