Invalid output of tf.image.transform operation when outputShape is different from input image shape (on cpu and wasm backends)

[andrusza2]

System information

• TensorFlow.js version: 3.18.0, master

Describe the current behavior
There is a bug in the implementation of the tf.image.transform operation in the cpu and wasm backends.
The problem manifests itself when we want to get the output of the operation with a shape other than the shape of the input image (using the outputShape parameter).

For example, using an input image of 256x256 and the outputShape set to [512,512], we get the following results:

• webgl-backend (correct output):

• cpu-backend/wasm-backend (incorrect output):

The problem also occurs when outputShape is smaller than the shape of the input image, for example, using 256x256 input image and setting outputShape to [96,96] we get the following results:

• webgl-backend - correct output, as above
• cpu-backend - incorrect output:

• wasm-backend - RuntimeError: memory access out of bounds

Describe the expected behavior

tf.image.transform operation works correctly regardless of the outputShape.

Standalone code to reproduce the issue

Snippet to reproduce:

// after setting backend to wasm - tf.setBackend('wasm')
img = tf.randomUniform([1,256,256,3])
proj_mat = tf.tensor([[2.,0.,0.,0.,2.,0.,0.,0.]])
tf.image.transform(img,proj_mat,'bilinear', 'constant', 0, [128,128])

Such a call causes a RuntimeError:

operation.ts:59 Uncaught RuntimeError: memory access out of bounds
    at tfjs-backend-wasm-simd.wasm:0x552b4
    at tfjs-backend-wasm-simd.wasm:0x553b0
    at tfjs-backend-wasm-simd.wasm:0x1937e
    at ccall (tfjs-backend-wasm.js:9:6302)
    at tfjs-backend-wasm.js:9:6692
    at Object.transform [as kernelFunc] (Transform.ts:95:3)
    at kernelFunc (engine.ts:644:22)
    at engine.ts:710:23
    at Engine.scopedRun (engine.ts:478:19)
    at Engine.runKernelFunc (engine.ts:706:10)

Other info / logs

I debugged this operation and found the source of the problem.
The error is related to the indexing of the output. The memory for the output tensor is allocated according to the outputShape, but when writing values to it - the strides values calculated for the input image (instead of the output size) are used to calculate the index:

const ind = b * batchStride + outY * rowStride + outX * colStride + channel;

rowStride and colStride values are incorrect here, which results in the output as in the examples above (and in the case of the outputShape smaller than the input shape - it even refers to the out of bounds memory).

The problematic code is here:

• cpu-backend
• wasm-backend

[andrusza2]

I created PR #6658 that should fix this problem.

[rthadur]

Thank you for the contribution.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No