import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
from tensorflow.keras import backend as K

BINSIZE = 1
XMIN = 0

def weighted_avg(inputs):
    # Calculate weighted sum of inputs
    ones = K.ones_like(inputs[0,:])       # [1, 1, 1, 1....]   (size Nouts)
    idx  = K.cumsum(ones)                 # [1, 2, 3, 4....]   (size Nouts)
    norm = K.sum(inputs, axis=-1, keepdims=True)          # normalization of all outputs by batch. shape is 1D array of size batch
    wavg = K.sum(idx*inputs, axis=-1, keepdims=True)/norm # array of size batch with weighted avg. of mean in units of bins
    wavg_cm = wavg*BINSIZE + XMIN         # array of size batch with weighted avg. of mean in physical units

    return wavg_cm

def make_model():
  inp = Input(shape=(4,))
  out = Lambda(weighted_avg)(inp)
  model = Model(inp, out)
  model.compile('adam', 'mse')
  return model

model = make_model()
model.summary()

import numpy as np
X = np.array([
    [1, 1, 1, 1],
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, 1]
])
model.predict(X)

array([[2.5],
       [1. ],
       [2. ],
       [3. ],
       [4. ]], dtype=float32)