Issue with tensorflow and output_transform

[pescap]

Hi,
The following code:

https://github.com/lululxvi/deepxde/blob/7322b23d1f74cfda62fcfb12bfa24d879e12f578/examples/pinn_forward/Helmholtz_Dirichlet_2d.py

freezes at step 0 with tensorflow for deepxde>1.2.
Everything works fine for deepxde=1.0,1.1,1.2.

I noticed it was due to the apply_output_transform(transform) function.

For example, If I set:

def transform(x, y):
    res = x[:, 0:1] * (1 - x[:, 0:1]) * x[:, 1:2] * (1 - x[:, 1:2])
    return y

it works fine, but if I add the x-terms part

def transform(x, y):
    res = x[:, 0:1] * (1 - x[:, 0:1]) * x[:, 1:2] * (1 - x[:, 1:2])
    return y * res

or return res it freezes. Any idea of what has to be changed?

[pescap]

This behavior does not appear for tf.compat.v1 or pytorch..

[ghost]

try dde.config.disable_xla_jit(). Maybe the problem caused by slice?

[pescap]

Thank you! It works fine now. To go further, how could we perform the slicing in a xla-friendly way?