grad returns non-deterministic values for some multi-dimentional inputs about catalyst

[grwlf]

Here is some additional information on this issue:

1. Disabling the bufferization inlining (as suggested by @erick-xanadu) makes the compiler output correct results in most cases. We put a Catalyst version with the inlining disabled as d934fa7 .
2. For this modified Catalyst, some tests always seem to be correct, but others are sometimes wrong.
3. An example of the flickering test is shown below. The fact that f returns multiple outputs seems to be important. We didn't see incorrect results from single-output versions of f.

import jax.numpy as jnp
from catalyst import qjit, grad as C_grad
from numpy.testing import assert_allclose

def f(x):
    return (x, jnp.stack([1*x, 2*x, 3*x]))

@qjit
def C_workflow():
    return C_grad(f, method='fd', argnum=[0])(jnp.zeros([4], dtype=float))

# Check the consistency
r0 = C_workflow()
for i in range(10):
    print(f"Iteration {i}")
    r1 = C_workflow()

    for a,b in zip(r0,r1):
        assert_allclose(a, b, rtol=1e-6, atol=1e-6)

An example output:

Iteration 0
Iteration 1
Iteration 2
Traceback (most recent call last):
  File "/workspace/src/vjp_study/issue1.py", line 19, in <module>
    assert_allclose(a, b, rtol=1e-6, atol=1e-6)
  File "/opt/venv/lib/python3.10/site-packages/numpy/testing/_private/utils.py", line 1527, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/opt/venv/lib/python3.10/site-packages/numpy/testing/_private/utils.py", line 844, in assert_array_compare
    raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=1e-06, atol=1e-06

Mismatched elements: 2 / 16 (12.5%)
Max absolute difference: 1.
Max relative difference: 2.13755206e+302
 x: array([[1., 0., 0., 0.],
       [0., 1., 0., 0.],
       [0., 0., 1., 0.],
       [0., 0., 0., 1.]])
 y: array([[4.678249e-303, 0.000000e+000, 0.000000e+000, 0.000000e+000],
       [4.678249e-303, 0.000000e+000, 0.000000e+000, 0.000000e+000],
       [4.678249e-303, 0.000000e+000, 1.000000e+000, 0.000000e+000],
       [4.678249e-303, 0.000000e+000, 0.000000e+000, 1.000000e+000]])

[erick-xanadu]

@grwlf can you show here the LLVM-IR?

[grwlf]

@grwlf can you show here the LLVM-IR?

Sure, this one? https://gist.github.com/grwlf/cf485111b55c19c3a2a9c6bbeed84777

[dime10]

Resolved by #121