isentropic · January 7, 2022 13:50 · May 5, 2020 · Nov 14, 2019 · Nov 14, 2019
diff --git a/histogram2d.py b/histogram2d.py
@@ -25,6 +25,11 @@ def get2dHistogram(x, 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)
     
     
     """

diff --git a/histogram2d.py b/histogram2d.py
@@ -3,7 +3,7 @@
 def get2dHistogram(x, y,
                    value_range,
                    nbins=100,
-                   dtype=dtypes.int32):
+                   dtype=tf.dtypes.int32):
     """
     Bins x, y coordinates of points onto simple square 2d histogram
     

diff --git a/histogram2d.py b/histogram2d.py
@@ -0,0 +1,38 @@
+import tensorflow as tf
+@tf.function
+def get2dHistogram(x, y,
+                   value_range,
+                   nbins=100,
+                   dtype=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.
+  
+    
+    
+    """
+    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!