Tensorflow histogram2d (Simple implementation)

histogram2d.py

@tf.function
def histogram(x, y,
              value_range,
              nbins=100,
              weights=None,
              bin_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.
    weights: The value to scale 
    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]

    if weights is None:
      hist_bins = tf.histogram_fixed_width_bins(y, y_range, nbins=nbins, dtype=bin_dtype)
      return tf.map_fn(lambda i: tf.histogram_fixed_width(x[hist_bins == i], x_range, nbins=nbins),
              tf.range(nbins))
    
    x_bins = tf.histogram_fixed_width_bins(x, x_range, nbins=nbins, dtype=bin_dtype)
    y_bins = tf.histogram_fixed_width_bins(y, y_range, nbins=nbins, dtype=bin_dtype)
    hist = tf.zeros((nbins, nbins), dtype=weights.dtype)
    indices = tf.transpose(tf.stack([y_bins, x_bins]))
    return tf.tensor_scatter_nd_add(hist, indices, weights)