tensorflow - Can you process a tensor in chunks in a custom Keras loss function? -

i trying write cusom keras loss function in process tensors in sub-vector chunks. example, if output tensor represented concatenation of quaternion coefficients (i.e. w,x,y,z,w,x,y,z...) might wish normalize each quaternion before calculating mean squared error in loss function like:

def norm_quat_mse(y_true, y_pred):     diff = y_pred - y_true     dist = 0     in range(0,16,4):         dist += k.sum( k.square(diff[i:i+4] / k.sqrt(k.sum(k.square(diff[i:i+4])))))      return dist/4 

while keras accept function without error , use in training, outputs different loss value when applied independent function , when using model.predict(), suspect not working properly. none of built-in keras loss functions use per-chunk processing approach, possible within keras' auto-differentiation framework?

try:

def norm_quat_mse(y_true, y_pred):     diff = y_pred - y_true     dist = 0     in range(0,16,4):         dist += k.sum( k.square(diff[:,i:i+4] / k.sqrt(k.sum(k.square(diff[:,i:i+4])))))      return dist/4 

you need know shape of y_true , y_pred (batch_size, output_size) need skip first dimension during computations.