Add contrib.rand_zipfian

python/mxnet/ndarray/contrib.py

@@ -18,9 +18,82 @@
 # coding: utf-8
 # pylint: disable=wildcard-import, unused-wildcard-import
 """Contrib NDArray API of MXNet."""
+import math
+from ..context import current_context
+from ..random import uniform
 try:
     from .gen_contrib import *
 except ImportError:
     pass
 
 __all__ = []
+
+def rand_log_uniform(true_classes, num_sampled, range_max, ctx=None):
+    """Draw random samples from an approximately log-uniform or Zipfian distribution.
+
+    This operation randomly samples *num_sampled* candidates the range of integers [0, range_max).
+    The elements of sampled_candidates are drawn with replacement from the base distribution.
+
+    The base distribution for this operation is an approximately log-uniform or Zipfian distribution:
+
+    P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1)
+
+    This sampler is useful when the true classes approximately follow such a distribution.
+    For example, if the classes represent words in a lexicon sorted in decreasing order of \
+    frequency. If your classes are not ordered by decreasing frequency, do not use this op.
+
+    Additionaly, it also returns the number of times each of the \
+    true classes and the sampled classes is expected to occur.
+
+    Parameters
+    ----------
+    true_classes : NDArray
+        A 1-D NDArray of the target classes.
+    num_sampled: int
+        The number of classes to randomly sample.
+    range_max: int
+        The number of possible classes.
+    ctx : Context
+        Device context of output. Default is current context. Overridden by
+        `mu.context` when `mu` is an NDArray.
+
+    Returns
+    -------
+    list of NDArrays
+        A 1-D `int64` `NDArray` for sampled candidate classes, a 1-D `float64` `NDArray` for \
+        the expected count for true classes, and a 1-D `float64` `NDArray` for the \
+        expected count for sampled classes.
+
+    Examples
+    --------
+    >>> true_cls = mx.nd.array[3]
+    >>> samples, exp_count_true, exp_count_sample = mx.nd.contrib.log_uniform(true_cls, 4, 5)
+    >>> samples
+    [[ 4.  0.]
+     [ 3.  2.]
+     [ 6.  2.]]
+    <NDArray 3x2 @cpu(0)>
+    >>> exp_count_true
+    [[ 4.  0.]
+     [ 3.  2.]
+     [ 6.  2.]]
+    <NDArray 3x2 @cpu(0)>
+    >>> exp_count_sample
+    [[ 4.  0.]
+     [ 3.  2.]
+     [ 6.  2.]]
+    <NDArray 3x2 @cpu(0)>
+    """
+    if ctx is None:
+        ctx = current_context()
+    log_range = math.log(range_max + 1)
+    rand = uniform(0, log_range, shape=(num_sampled,), dtype='float64', ctx=ctx)
+    # make sure sampled_classes are in the range of [0, range_max)
+    sampled_classes = (rand.exp() - 1).astype('int64') % range_max
+
+    true_classes = true_classes.as_in_context(ctx).astype('float64')
+    expected_count_true = ((true_classes + 2.0) / (true_classes + 1.0)).log() / log_range
+    # cast sampled classes to fp64 to avoid interget division
+    sampled_cls_fp64 = sampled_classes.astype('float64')
+    expected_count_sampled = ((sampled_cls_fp64 + 2.0) / (sampled_cls_fp64 + 1.0)).log() / log_range
+    return [sampled_classes, expected_count_true, expected_count_sampled]

python/mxnet/ndarray/random.py

@@ -17,6 +17,7 @@
 
 """Random distribution generator NDArray API of MXNet."""
 
+import math
 from ..base import numeric_types, _Null
 from ..context import current_context
 from . import _internal

python/mxnet/symbol/contrib.py

@@ -18,9 +18,84 @@
 # coding: utf-8
 # pylint: disable=wildcard-import, unused-wildcard-import
 """Contrib Symbol API of MXNet."""
+import math
+from ..context import current_context
+from .random import uniform
+from .symbol import Symbol
 try:
     from .gen_contrib import *
 except ImportError:
     pass
 
 __all__ = []
+
+def rand_log_uniform(true_classes, num_sampled, range_max, ctx=None):
+    """Draw random samples from an approximately log-uniform or Zipfian distribution.
+
+    This operation randomly samples *num_sampled* candidates the range of integers [0, range_max).
+    The elements of sampled_candidates are drawn with replacement from the base distribution.
+
+    The base distribution for this operation is an approximately log-uniform or Zipfian distribution:
+
+    P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1)
+
+    This sampler is useful when the true classes approximately follow such a distribution.
+    For example, if the classes represent words in a lexicon sorted in decreasing order of \
+    frequency. If your classes are not ordered by decreasing frequency, do not use this op.
+
+    Additionaly, it also returns the number of times each of the \
+    true classes and the sampled classes is expected to occur.
+
+    Parameters
+    ----------
+    true_classes : Symbol
+        The target classes in 1-D.
+    num_sampled: int
+        The number of classes to randomly sample.
+    range_max: int
+        The number of possible classes.
+    ctx : Context
+        Device context of output. Default is current context. Overridden by
+        `mu.context` when `mu` is an NDArray.
+
+    Returns
+    -------
+    list of Symbols
+        A 1-D `int64` `Symbol` for sampled candidate classes, a 1-D `float64` `Symbol` for \
+        the expected count for true classes, and a 1-D `float64` `Symbol` for the \
+        expected count for sampled classes.
+
+    Examples
+    --------
+    >>> true_cls = mx.nd.array[3]
+    >>> samples, exp_count_true, exp_count_sample = mx.nd.contrib.log_uniform(true_cls, 4, 5)
+    >>> samples
+    [[ 4.  0.]
+     [ 3.  2.]
+     [ 6.  2.]]
+    <NDArray 3x2 @cpu(0)>
+    >>> exp_count_true
+    [[ 4.  0.]
+     [ 3.  2.]
+     [ 6.  2.]]
+    <NDArray 3x2 @cpu(0)>
+    >>> exp_count_sample
+    [[ 4.  0.]
+     [ 3.  2.]
+     [ 6.  2.]]
+    <NDArray 3x2 @cpu(0)>
+    """
+    assert(isinstance(true_classes, Symbol)), "unexpected type %s" % type(true_classes)
+    if ctx is None:
+        ctx = current_context()
+    log_range = math.log(range_max + 1)
+    rand = uniform(0, log_range, shape=(num_sampled,), dtype='float64', ctx=ctx)
+    # make sure sampled_classes are in the range of [0, range_max)
+    sampled_classes = (rand.exp() - 1).astype('int64') % range_max
+
+    true_classes = true_classes.astype('float64')
+    expected_count_true = ((true_classes + 2.0) / (true_classes + 1.0)).log() / log_range
+    # cast sampled classes to fp64 to avoid interget division
+    sampled_cls_fp64 = sampled_classes.astype('float64')
+    expected_count_sampled = ((sampled_cls_fp64 + 2.0) / (sampled_cls_fp64 + 1.0)).log() / log_range
+    return [sampled_classes, expected_count_true, expected_count_sampled]

src/operator/tensor/elemwise_unary_op.h

@@ -357,6 +357,7 @@ struct CastParam : public dmlc::Parameter<CastParam> {
     .add_enum("float16", mshadow::kFloat16)
     .add_enum("uint8", mshadow::kUint8)
     .add_enum("int32", mshadow::kInt32)
+    .add_enum("int64", mshadow::kInt64)
     .describe("Output data type.");
   }
 };

tests/python/unittest/test_random.py

@@ -412,6 +412,36 @@ def test_multinomial_generator():
                  for _ in range(10)])
         verify_generator(generator=generator_mx_same_seed, buckets=buckets, probs=probs)
 
+def test_log_uniform_generator():
+    mx.random.seed(0)
+    # dummy true classes
+    num_true = 5
+    num_sampled = 1000
+    range_max = 20
+
+    def compute_expected_counts():
+        # P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1)
+        classes = mx.nd.arange(0, range_max)
+        expected_counts = ((classes + 2).log() - (classes + 1).log()) / np.log(range_max + 1)
+        return expected_counts
+
+    exp_cnt = compute_expected_counts()
+
+    # test ndarray
+    true_classes = mx.nd.random.uniform(0, range_max, shape=(num_true,)).astype('int32')
+    sampled_classes, exp_cnt_true, exp_cnt_sampled = mx.nd.contrib.rand_log_uniform(true_classes, num_sampled, range_max)
+    mx.test_utils.assert_almost_equal(exp_cnt_sampled.asnumpy(), exp_cnt[sampled_classes].asnumpy(), rtol=1e-1, atol=1e-2)
+    mx.test_utils.assert_almost_equal(exp_cnt_true.asnumpy(), exp_cnt[true_classes].asnumpy(), rtol=1e-1, atol=1e-2)
+
+    # test symbol
+    true_classes_var = mx.sym.var('true_classes')
+    outputs = mx.sym.contrib.rand_log_uniform(true_classes_var, num_sampled, range_max)
+    outputs = mx.sym.Group(outputs)
+    executor = outputs.bind(mx.context.current_context(), {'true_classes' : true_classes})
+    executor.forward()
+    sampled_classes, exp_cnt_true, exp_cnt_sampled = executor.outputs
+    mx.test_utils.assert_almost_equal(exp_cnt_sampled.asnumpy(), exp_cnt[sampled_classes].asnumpy(), rtol=1e-1, atol=1e-2)
+    mx.test_utils.assert_almost_equal(exp_cnt_true.asnumpy(), exp_cnt[true_classes].asnumpy(), rtol=1e-1, atol=1e-2)
 
 if __name__ == '__main__':
     import nose