Histograma 2D con pesos en Tensorflow

Necesito usar un histograma en 2D con pesos en Tensorflow. Sin embargo, la única solución que encontré en aquí no implementa el argumento de pesos como en la implementación de numpy. Al analizar el código fuente de la implementación de numpy, pude hacer lo siguiente en la implementación de tensorflow encontrada, pero aún no funciona como se espera:

def get2dHistogram(x, y, weights,
               value_range,
               nbins=100,
               dtype=tf.dtypes.int32):

x_range = value_range[0]
y_range = value_range[1]

x = tf.histogram_fixed_width_bins(x, y_range, nbins=tf.size(x), dtype=dtype)
x = tf.math.bincount(x, weights=weights, minlength=tf.size(y))
y = tf.histogram_fixed_width_bins(y, y_range, nbins=tf.size(y), dtype=dtype)
y = tf.math.bincount(y, weights=weights, minlength=tf.size(y))

histy_bins = tf.histogram_fixed_width_bins(y, y_range, nbins=nbins, dtype=dtype)

H = tf.map_fn(lambda i: tf.histogram_fixed_width(x[histy_bins == i], x_range, nbins=nbins), tf.range(nbins))
return H # Matrix!

Confieso que no sé exactamente cómo se implementan los pesos en la versión de numpy, pero esa fue mi mejor suposición hasta ahora. ¿Alguien puede ayudarme?