tf.einsum appears incorrect for some cases

[zachlipp]

Hey there, I am playing around with tf.js for the first time and I have noticed some discrepancies from the results I get using einsum here versus in Python. I may be missing something, but I searched the issues and PRs on this repo and didn't find anything glaring so I'm posting here.

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): MacOS Sonoma 14.3 (23D56)
• TensorFlow.js installed from (npm or script link): First noticed with the latest release, but reproducible from the tf.js interactive examples with some edits
• TensorFlow.js version (use command below): latest
• Browser version: Produced with Firefox 122.0.1 and Safari Version 17.3
• Tensorflow.js Converter Version: I do not know how to find this

Describe the current behavior
Per the docs, tf.einsum supports Einstein summation notation. There are some documented unsupported cases (e.g. more than 2 tensors), but I expect that within those constraints, the results I get should match the results from other implementations (for example, Tensorflow Python).

Examples:

Setup

JS:

const x = tf.tensor2d([[1, 2, 3], [4, 5, 6]]);
const y = tf.tensor2d([[0, 1], [2, 3], [4, 5]]);

Python:

x = tf.constant([[1, 2, 3], [4, 5, 6]]);
y = tf.constant([[0, 1], [2, 3], [4, 5]]);

1

tf.einsum('ij,jk->k', x, y)

returns [[0, 14, 36], [5, 21, 45]] in JS, but in Python this yields [50, 71]

2

tf.einsum('ij,jk->j', x, y)

returns [[0 , 21 ], [42, 63 ], [84, 105]] in JS but in Python this yields [5, 35, 81]

3

tf.einsum('ij,jk->', x, y)

returns [50, 71] in JS, but in Python this returns 121.

Link to Python examples.

I'm new to JS, so I don't know the preferred way to link examples. However, this data is from the first example of the tf.js einsum documentation. Modifying the einsum notation there can reproduce these results.

[gaikwadrahul8]

Hi, @zachlipp

Thank you for bringing this issue to our attention and I was trying to replicate the same issue from my end also and I'm also getting same results which you mentioned in your issue template as far my current understanding it's happening because of tensor shapes TensorFlow Python and TensorFlow.js handle certain tensor shapes differently in tf.einsum. Specifically, when summing over multiple indices to produce a scalar output, Python reduces across those dimensions, while JavaScript may produce a tensor with the remaining dimensions.

TensorFlow Python and TensorFlow.js handle dimension reduction differently in tf.einsum when summing over multiple axes due to which there is discrepancy between the output of tf.einsum in TensorFlow Python and TensorFlow.js if I'm not wrong, please correct me if I'm wrong

For reference I have added output log below and Google colab notebook :

Thank you for your understanding and patience.

[zachlipp]

Hey there! I am happy to hear you can reproduce my results. it definitely appears that the js implementation will have an "extra" axis every so often. I guess my question now is: Do you / the team consider this discrepancy between the Python and JS versions of this API a bug? Researchers often use this notation in Python, and I was hoping that tf.js could help me reproduce those results in JS.