issue with tf.norm

[MYOUSSF]

I have the following implementation of a PDE

def pde(t, y):
r = tf.norm(y, axis=1, keepdims=True)
y0_tt = dde.grad.hessian(y, t, component=0)
y1_tt = dde.grad.hessian(y, t, component=1)
return [y0_tt + mu * y[:,0:1] / r3, y1_tt + mu * y[:,1:2] / r3]

Here mu is a constant. The rest of the code is standard and irrelevant to the issue as the only part that is problematic is the division by the norm of y: r = tf.norm(y, axis=1, keepdims=True).

The shape of y is (none, 2). The shapes of y[:,0:1] and r is (none ,1), so it makes sense to do y[:,0:1] / r**3. However this is the source of the problem. If, for example, I take r to be a scalar or remove it altogether from the code, everything works fine. But with r defined as above, I get the following errors:

Training model...

0 [nan, nan, 6.40e+07, 0.00e+00, 1.72e-02, 4.17e+01] [nan, nan, 6.40e+07, 0.00e+00, 1.72e-02, 4.17e+01] []

Best model at step 0:
train loss: inf
test loss: inf
test metric:

'train' took 1.271964 s

Saving loss history to /content/loss.dat ...
Saving training data to /content/train.dat ...
Saving test data to /content/test.dat ...

ValueError Traceback (most recent call last)
in <cell line: 62>()
60
61 losshistory, train_state = model.train(iterations=4000)
---> 62 dde.saveplot(losshistory, train_state, issave=True, isplot=True)

4 frames
/usr/local/lib/python3.10/dist-packages/deepxde/utils/external.py in saveplot(loss_history, train_state, issave, isplot, loss_fname, train_fname, test_fname, output_dir)
173 test_fname = os.path.join(output_dir, test_fname)
174 save_loss_history(loss_history, loss_fname)
--> 175 save_best_state(train_state, train_fname, test_fname)
176
177 if isplot:

/usr/local/lib/python3.10/dist-packages/deepxde/utils/external.py in save_best_state(train_state, fname_train, fname_test)
342 print("Saving test data to {} ...".format(fname_test))
343 if y_test is None:
--> 344 test = np.hstack((train_state.X_test, best_y))
345 if best_ystd is None:
346 np.savetxt(fname_test, test, header="x, y_pred")

/usr/local/lib/python3.10/dist-packages/numpy/core/overrides.py in hstack(*args, **kwargs)

/usr/local/lib/python3.10/dist-packages/numpy/core/shape_base.py in hstack(tup)
343 return _nx.concatenate(arrs, 0)
344 else:
--> 345 return _nx.concatenate(arrs, 1)
346
347

/usr/local/lib/python3.10/dist-packages/numpy/core/overrides.py in concatenate(*args, **kwargs)

ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s)

[MYOUSSF]

I don't know if this is an issue with DeepXDE or I'm missing something. Any help is appreciated.

[MYOUSSF]

I found the cause of the problem. It's division by zero.