Tensorflow histogram2d (Simple implementation)

histogram2d.py

import tensorflow as tf
@tf.function
def get2dHistogram(x, y,
                   value_range,
                   nbins=100,
                   dtype=tf.dtypes.int32):
    """
    Bins x, y coordinates of points onto simple square 2d histogram
    
    Given the tensor x and y:
    x: x coordinates of points
    y: y coordinates of points
    this operation returns a rank 2 `Tensor` 
    representing the indices of a histogram into which each element
    of `values` would be binned. The bins are equal width and
    determined by the arguments `value_range` and `nbins`.
    
    
  Args:
    x:  Numeric `Tensor`.
    y: Numeric `Tensor`.
    value_range[0] lims for x
    value_range[1] lims for y
    
    nbins:  Scalar `int32 Tensor`.  Number of histogram bins.
    dtype:  dtype for returned histogram.
  
  Example:
  N = 1000
  xs =  tf.random.normal([N])
  ys =  tf.random.normal([N])
  get2dHistogram(xs, ys, ([-5.0, 5.0], [-5.0, 5.0]),  50)
    
    
    """
    x_range = value_range[0]
    y_range = value_range[1]

    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!

I managed to make it work using the tensorflow-probability package for the histogram implementation. The final is below:

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

x_range = tf.linspace(value_range[0][0], value_range[0][1], num=nbins+1)
x = tf.clip_by_value(x, value_range[0][0], value_range[0][1])

histy_bins = tf.histogram_fixed_width_bins(y, value_range[1], nbins=nbins, dtype=dtype)

hists = []
for i in range(nbins):
_x = x[histy_bins == i]
_w = weights[histy_bins == i]
hist = tfp.stats.histogram(_x, edges=x_range, weights=_w)
hists.append(hist)

return tf.stack(hists, axis=0)`