title: Equivalent of numpy.digitize in tensorflow
body: <p>I am working on a customised loss function that uses <code>numpy.digitize()</code> internally. The loss is minimised for a set of parameters that are the <code>bins</code> values used in digitize method. In order to use the <code>tensorflow</code> optimisers, I would like to know if there an equivalent implementation of <code>digitize</code> in <code>tensorflow</code>? if not is there a good way to implement a workaround? </p>
  
  <p>Here a numpy version:</p>
  
  <pre class="lang-py prettyprint-override"><code>def fom_func(b, n):
      np.where((b &gt; 0) &amp; (n &gt; 0), np.sqrt(2*(n*np.log(np.divide(n,b)) + b - n)),0) 
  def loss(param, X, y):
      param = np.sort(np.asarray(param))
      nbins = param.shape[0]
      score = 0
      y_pred = np.digitize(X, param)
      for c in np.arange(nbins):
          b = np.where((y==0) &amp; (y_pred==c), 1, 0).sum()
          n = np.where((y_pred==c), 1, 0).sum()
          score += fom_func(b,n)**2
      return -np.sqrt(score)
  </code></pre>