tape.gradient throws InvalidArgumentError: Input to reshape is a tensor with 1 values, but the requested shape has 400 [Op:Reshape]

Question

When implementing an unsupervised learning algorithm, that has the goal of inferring a set of coordinates which follow a given angular two-point correlation function, I get an error with my custom gradient descent algorithm.

The loss function is basically the difference between the angular_tpcf of the output coordinates of the network and the goal correlation function. For this calculation, I need a Numpy function which I currently cannot translate into a TensorFlow function, I use the tf.py_function wrapper instead. When calculating the gradient of the loss, I get a strange error pointing to a faulty reshape of the gradient. The error is:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
/var/folders/9n/5rsrx6tx2jx4234dz2mgmqhc0000gn/T/ipykernel_27441/710230524.py in <module>
     38              loss = tpcfloss(coords=yout, corgoal=corgoal)
     39              #print(loss, model.trainable_variables)
---> 40         grad = tape.gradient(loss, model.trainable_variables) # Backpropagation
     41         print(grad)
     42         optimizer.apply_gradients(zip(grad, model.trainable_variables)) # Update network weights

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/eager/backprop.py in gradient(self, target, sources, output_gradients, unconnected_gradients)
   1082                           for x in nest.flatten(output_gradients)]
   1083 
-> 1084     flat_grad = imperative_grad.imperative_grad(
   1085         self._tape,
   1086         flat_targets,

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/eager/imperative_grad.py in imperative_grad(tape, target, sources, output_gradients, sources_raw, unconnected_gradients)
     69         "Unknown value for unconnected_gradients: %r" % unconnected_gradients)
     70 
---> 71   return pywrap_tfe.TFE_Py_TapeGradient(
     72       tape._tape,  # pylint: disable=protected-access
     73       target,

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/eager/backprop.py in _gradient_function(op_name, attr_tuple, num_inputs, inputs, outputs, out_grads, skip_input_indices, forward_pass_name_scope)
    157       gradient_name_scope += forward_pass_name_scope + "/"
    158     with ops.name_scope(gradient_name_scope):
--> 159       return grad_fn(mock_op, *out_grads)
    160   else:
    161     return grad_fn(mock_op, *out_grads)

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/ops/array_grad.py in _ReshapeGrad(op, grad)
    781 def _ReshapeGrad(op, grad):
    782   return [
--> 783       array_ops.reshape(
    784           _IndexedSlicesToTensorNoWarning(grad), array_ops.shape(op.inputs[0])),
    785       None

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    204     """Call target, and fall back on dispatchers if there is a TypeError."""
    205     try:
--> 206       return target(*args, **kwargs)
    207     except (TypeError, ValueError):
    208       # Note: convert_to_eager_tensor currently raises a ValueError, not a

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/ops/array_ops.py in reshape(tensor, shape, name)
    194     A `Tensor`. Has the same type as `tensor`.
    195   """
--> 196   result = gen_array_ops.reshape(tensor, shape, name)
    197   tensor_util.maybe_set_static_shape(result, shape)
    198   return result

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/ops/gen_array_ops.py in reshape(tensor, shape, name)
   8392       return _result
   8393     except _core._NotOkStatusException as e:
-> 8394       _ops.raise_from_not_ok_status(e, name)
   8395     except _core._FallbackException:
   8396       pass

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6939   message = e.message + (" name: " + name if name is not None else "")
   6940   # pylint: disable=protected-access
-> 6941   six.raise_from(core._status_to_exception(e.code, message), None)
   6942   # pylint: enable=protected-access
   6943 

/opt/homebrew/anaconda3/envs/train/lib/python3.9/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Input to reshape is a tensor with 1 values, but the requested shape has 400 [Op:Reshape]

The Code for the network is:

#definition of the model
def define_model(outputdim):
    inputlay = keras.layers.Input(shape=(1000,))
    #some inputlayers
    x =  keras.layers.Dense(128,activation='relu')(inputlay)
    x = keras.layers.Dense(128,activation='relu')(x)
    x = keras.layers.Dense(2*outputdim,activation='relu')(x)
    output = keras.layers.Reshape((outputdim,2))(x)
    return keras.Model(inputlay,output)

#some parameters
nclus = 200
n_epochs = 10
n_batches = 10

randomsforwthetadec = 90*np.random.random((nclus,))
randomsforwthetara = 180*np.random.random((nclus,))
randomsforwt = np.vstack((randomsforwthetara,randomsforwthetadec)).T

thetabinning = np.linspace(0,10,36)

model = define_model(nclus)

#definition of the loss function
def tpcfloss(coords,corgoal):
    #coords = coords.numpy()
    return keras.backend.sum(keras.backend.abs(tf.py_function(hmock.angular_tpcf,[coords,thetabinning,randomsforwt],tf.float64)-tf.convert_to_tensor(corgoal[1:36])))


#setting up training
optimizer = tf.keras.optimizers.Adam()
loss_train = np.zeros(shape=(n_epochs,), dtype=np.float32)
loss_val = np.zeros(shape=(n_epochs,))

#training loop
for epoch in range(n_epochs):
    epoch_loss_avg = tf.keras.metrics.Mean() # Keeping track of the training loss
    print('==== Epoch #{0:3d} ===='.format(epoch))
    for batch in range(n_batches):
        xin = np.random.random((1,1000))
        with tf.GradientTape() as tape: # Forward pass
             yout = model(xin, training=True)
             #print(yout)
             loss = tpcfloss(coords=yout, corgoal=corgoal)
             #print(loss, model.trainable_variables)
        grad = tape.gradient(loss, model.trainable_variables) # Backpropagation
        print(grad)
        optimizer.apply_gradients(zip(grad, model.trainable_variables)) # Update network weights
    epoch_loss_avg(loss)
    loss_train[epoch] = epoch_loss_avg.result()
    print('---- Training ----')
    print('Loss  =  {0:.3f}'.format(loss_train[epoch]))

Thank you in advance to any and all answers!