Custom loss function with tf.function produces unexpected test outcome

[cyan-at]

We are trying to shape a pinn output deepxde.IC.PointSetBC into a pdf (sum = 1, non-negative)

Here is the BC:

np.transpose(np.linspace(0, 6, N))
time_0=np.hstack((x_T,T_0*np.ones((len(x_T), 1))))
rho_0=pdf1d_0(x_T).reshape(len(x_T),1)
rho_0 = np.where(rho_0 < 0, 0, rho_0)
rho_0 = rho_0 / np.sum(np.abs(rho_0))
rho_0_BC = dde.icbc.PointSetBC(time_0, rho_0, component=1)

I wanted to test what the pinn produces for a few simple loss functions:

def Z(y_true, y_pred):
    return  tf.abs(**target** - tf.reduce_sum(y_pred))

def Z2(y_true, y_pred):
    loss = tf.py_function(
        func=Z,
        inp=[y_true, y_pred],
        Tout=tf.float32
    )
    return loss

Z is a simple loss to steer the pinn to shaping the BC to have a sum == target

When target above == 0, the output has a sum of 6
When target above == 0.5, the output has a sum of 9
When target above == 1, output has sum of 12

The pattern seems to be: output_sum = state_max + state_max * target

And indeed, if I want the output_sum == 1, I set target == 5/6 and it does produce what I expect

I was wondering why this is the case? Why does the output's sum not equal what the loss function is penalizing? How might it be related to state_max - state_min?

Thank you

[cyan-at]

See attached for the example outputs for the various target values

target == 0, output == 6

target == 0.5, output == 9

target == 1, output == 12

[lululxvi]

No idea. Probably due to tf.py_function? I have no experience with tf.py_function.