Add discrete distributions to dpnp.random.Generator

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added `dpnp.top_k` to return the `k` largest or smallest elements and their indices along an axis, with a `sorted` keyword to return them in sorted order [#3133](https://github.com/IntelPython/dpnp/pull/3133)
 * Added the `dpnp.random.Generator` API with `default_rng`, `BitGenerator`, `SeedSequence` and the `Philox4x3210` bit generator, sampling on the device with the oneMKL device RNG API [#3157](https://github.com/IntelPython/dpnp/pull/3157)
 * Added the continuous distributions of `numpy.random.Generator`, such as `gamma`, `beta` and `vonmises`, to `dpnp.random.Generator` [#3158](https://github.com/IntelPython/dpnp/pull/3158)
+* Added `integers` and the discrete distributions of `numpy.random.Generator`, such as `poisson`, `binomial` and `hypergeometric`, to `dpnp.random.Generator`, along with `noncentral_chisquare` and `noncentral_f` [#3159](https://github.com/IntelPython/dpnp/pull/3159)
 
 ### Changed
 

doc/reference/random.rst

@@ -17,19 +17,21 @@ host-side oneMKL RNG.
 
     rng = dpnp.random.default_rng(seed=42)
     x = rng.standard_normal((3, 4))                # on the default device
-    y = rng.random(5, usm_type="shared")
+    y = rng.integers(0, 10, size=5, usm_type="shared")
 
     rng = dpnp.random.default_rng(42, device="cpu")
 
 A generator draws all its samples on one SYCL queue, chosen with the
 ``device`` or ``sycl_queue`` keyword when it is created, and the samples
-depend only on the seed, not on the device. All samples are
+depend only on the seed, not on the device, except that the discrete
+distributions also depend on its support of double precision. All samples are
 :class:`dpnp.ndarray`, zero-dimensional when ``size`` is ``None``.
 
 Differences from NumPy
 ----------------------
 
-* The same seed gives the same stream on every device.
+* The same seed gives the same stream on every device, but for the discrete
+  distributions noted below.
   The streams differ from NumPy's.
 * The floating-point distributions take a ``dtype`` keyword, as in CuPy,
   which defaults to the default floating-point type of the device. All
@@ -39,7 +41,20 @@ Differences from NumPy
   not, since that would wait for the device.
   Invalid ones give undefined values, such as NaN or the smallest ``int64``,
   but the kernels handle them safely: they never loop without end or access
-  memory out of bounds.
+  memory out of bounds. Floating-point arrays given for integer parameters
+  are converted as by :meth:`dpnp.ndarray.astype`, whose result for NaN and
+  out-of-range values is undefined.
+* The discrete distributions return ``int64`` values. They compute in double
+  precision on devices supporting it, and in single precision otherwise.
+  There, the streams differ from the ones of devices with double precision,
+  and some samples are only multiples of a power of two: the ones above
+  ``2**24`` of ``geometric``, ``logseries`` and ``zipf``, and the ones of
+  ``poisson`` and ``binomial`` with a standard deviation above about
+  ``2**22``.
+* ``zipf`` follows the Zipf distribution truncated to the largest ``int64``.
+  NumPy's rejection test rounds to accepting every sample above ``2**53``,
+  which makes its tail too heavy for ``a`` close to 1.
+* ``hypergeometric`` requires ``ngood`` and ``nbad`` below ``10**9``.
 
 Random Generator
 ----------------
@@ -58,6 +73,7 @@ Simple random data
    :toctree: generated/
    :nosignatures:
 
+   Generator.integers
    Generator.random
 
 Distributions
@@ -68,16 +84,24 @@ Distributions
    :nosignatures:
 
    Generator.beta
+   Generator.binomial
    Generator.chisquare
    Generator.exponential
    Generator.f
    Generator.gamma
+   Generator.geometric
    Generator.gumbel
+   Generator.hypergeometric
    Generator.laplace
    Generator.logistic
    Generator.lognormal
+   Generator.logseries
+   Generator.negative_binomial
+   Generator.noncentral_chisquare
+   Generator.noncentral_f
    Generator.normal
    Generator.pareto
+   Generator.poisson
    Generator.power
    Generator.rayleigh
    Generator.standard_cauchy
@@ -90,6 +114,7 @@ Distributions
    Generator.vonmises
    Generator.wald
    Generator.weibull
+   Generator.zipf
 
 Bit generators
 --------------

dpnp/AGENTS.md

@@ -21,7 +21,7 @@
 - Extension calls are ordered with `_manager = dpctl.utils.SequentialOrderManager[exec_q]`: pass `depends=_manager.submitted_events`, then `_manager.add_event_pair(ht_ev, ev)`. Don't add blocking waits.
 - Use `dpnp.default_float_type(...)` or `map_dtype_to_device(...)` for the default float type.
 - Deprecated ufuncs use `DPNPDeprecatedUnaryFunc(..., deprecated_msg=...)`.
-- A `dpnp.random` distribution is a `Generator` method: it calls `self._check_queues(...)` on its array parameters (raises `ExecutionPlacementError`), checks host parameters as NumPy does, and draws through `_cont` or `_float_sample` into `BitGenerator._fill`. Parameters on the device are not checked, as the check waits for the device. Temporaries take the output's `usm_type`.
+- A `dpnp.random` distribution is a `Generator` method: it calls `self._check_queues(...)` on its array parameters (raises `ExecutionPlacementError`), checks host parameters as NumPy does, and draws through `_cont`, `_float_sample` or `_discrete_sample` into `BitGenerator._fill`. Parameters on the device are not checked, as the check waits for the device. Temporaries take the output's `usm_type`.
 - Unsupported NumPy keywords raise `NotImplementedError` (`dpnp.check_limitations`) and are listed under `Limitations` in the docstring.
 - The Array API spec version is set in three places that change together: `__array_api_version__` in `dpnp/tensor/_array_api.py`, `ARRAY_API_TESTS_VERSION` in `.github/workflows/conda-package.yml` and `doc/reference/array_api.rst` (plus the version test in `dpnp/tests/test_ndarray.py`). Inspection utilities come from `__array_namespace_info__` in `dpnp_array_api_info.py`.
 - `dpnp.tensor` replaces `dpctl.tensor`, which no longer exists here; `SyclQueue`, `SyclDevice`, `dpctl.memory` and `SequentialOrderManager` still come from `dpctl`.

dpnp/libtensor/AGENTS.md

@@ -23,7 +23,7 @@ A new extension's `CMakeLists.txt` lists its sources and calls `build_dpnp_pybin
 - Guard double, complex<double> and half with `dpnp::tensor::type_utils::validate_type_for_device<T>(q)`; don't use double literals or casts in kernels that must run without fp64.
 - `blas`, `fft`, `lapack` and `sparse` build against either oneMKL or oneMath (`USE_ONEMATH`); `vm` has kernels only with oneMKL: with oneMath it builds empty and `_is_available()` returns `False`; `rng` uses the header-only oneMKL device RNG API and links no oneMKL library.
 - Release the GIL around blocking oneMKL calls (`py::gil_scoped_release`).
-- `rng` values depend only on the engine state and the number of elements, never on the device or the layout chosen for it (`extensions/rng/kernels/sample.hpp`); each engine's `items_per_work_item` (`kernels/engines.hpp`) is part of its stream definition. Parameters on the device are unchecked, so a rejection sampler returns NaN for the inputs that would make it loop forever, including values that are valid but round badly in fp32.
+- `rng` values depend only on the engine state and the number of elements, never on the device or the layout chosen for it (`extensions/rng/kernels/sample.hpp`); each engine's `items_per_work_item` (`kernels/engines.hpp`) is part of its stream definition. Parameters on the device are unchecked, so a rejection sampler returns NaN or the smallest `int64` for the inputs that would make it loop forever, including values that are valid but round badly in fp32. Use `umul_hi` from `kernels/umul_hi.hpp` instead of `sycl::mul_hi`, which DPC++ 2026.2 miscompiles on Intel GPUs.
 - An extension whose kernels may be launched over more than `INT_MAX` work-items passes `LARGE_RANGES` to `build_dpnp_pybind11_ext`, which adds `-fno-sycl-id-queries-fit-in-int` at compile and link time (DPC++ 2026.2 and later need it at link time). `source/CMakeLists.txt` and `dpnp_backend_c` pass it to both `target_compile_options` and `target_link_options`.
 - clang-format 23 formats C++ (`.clang-format`: LLVM base, 4-space indent, 80 columns); gersemi formats CMake at 88 columns, except `dpnp/backend/cmake/Modules/`.
 

dpnp/libtensor/extensions/rng/CMakeLists.txt

@@ -29,9 +29,14 @@
 
 set(python_module_name _rng_impl)
 set(_module_src
+    ${CMAKE_CURRENT_SOURCE_DIR}/binomial.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/discrete.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/exponential.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/gamma.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/integers.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/noncentral.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/normal.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/poisson.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/rng_py.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/sample.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/uniform.cpp

dpnp/libtensor/extensions/rng/binomial.cpp

@@ -0,0 +1,44 @@
+//*****************************************************************************
+// Copyright (c) 2026, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+// - Neither the name of the copyright holder nor the names of its contributors
+//   may be used to endorse or promote products derived from this software
+//   without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+#include <pybind11/pybind11.h>
+
+#include "common.hpp"
+#include "distributions.hpp"
+
+namespace dpnp::extensions::rng
+{
+static discrete_dispatch_table_t binomial_dispatch_table;
+
+void init_binomial(py::module_ m)
+{
+    def_discrete_sampler<rng_ns::BinomialSampler>(
+        m, "_binomial",
+        "Fill the result with binomial samples of n trials of probability p",
+        binomial_dispatch_table);
+}
+} // namespace dpnp::extensions::rng

dpnp/libtensor/extensions/rng/common.hpp

@@ -81,6 +81,10 @@ typedef sycl::event (*sample_fn_ptr_t)(sycl::queue &,
 using sample_dispatch_table_t =
     sample_fn_ptr_t[rng_ns::num_engines][td_ns::num_types];
 
+// Discrete samplers, whose result is always int64, indexed by the engine and
+// by whether they compute in double (1) or in float (0)
+using discrete_dispatch_table_t = sample_fn_ptr_t[rng_ns::num_engines][2];
+
 namespace detail
 {
 template <typename T>
@@ -138,8 +142,11 @@ sycl::event submit_sample(sycl::queue &exec_q,
 } // namespace detail
 
 // Checks the parameters and the result against SamplerT, then submits the
-// kernel, unless there is nothing to sample.
-template <typename EngineTraitsT, typename SamplerT>
+// kernel, unless there is nothing to sample. ComputeT is the type the
+// sampler computes in, which the device must support.
+template <typename EngineTraitsT,
+          typename SamplerT,
+          typename ComputeT = typename SamplerT::result_type>
 sycl::event sample_impl(sycl::queue &exec_q,
                         const rng_ns::EngineState &state,
                         const std::size_t nelems,
@@ -162,6 +169,7 @@ sycl::event sample_impl(sycl::queue &exec_q,
     }
 
     detail::validate_type_for_device<T>(exec_q);
+    detail::validate_type_for_device<ComputeT>(exec_q);
     detail::validate_types_for_device<param_types>(
         exec_q, std::make_index_sequence<n_params>{});
 
@@ -195,6 +203,12 @@ struct RealTypes
 {
 };
 
+// Types the integer samplers support
+template <typename T>
+struct IntegerTypes : std::bool_constant<std::is_integral_v<T>>
+{
+};
+
 // The raw output of a bit generator, whole or as 32-bit halves
 template <typename T>
 struct RawTypes : std::bool_constant<std::is_same_v<T, std::uint64_t> ||
@@ -238,11 +252,27 @@ void init_sample_dispatch_table(sample_dispatch_table_t &table)
         table[philox]);
 }
 
+template <template <typename> class SamplerT>
+void init_discrete_dispatch_table(discrete_dispatch_table_t &table)
+{
+    static_assert(rng_ns::num_engines == 1,
+                  "Every engine must have a row in the dispatch table");
+
+    using Philox = rng_ns::Philox4x32x10;
+    constexpr int philox = static_cast<int>(rng_ns::engine_id::philox4x32x10);
+    table[philox][0] = sample_impl<Philox, SamplerT<float>, float>;
+    table[philox][1] = sample_impl<Philox, SamplerT<double>, double>;
+}
+
 // Looks up the sampler of `result`'s type, and throws if there is none
 extern sample_fn_ptr_t lookup_sampler(const sample_dispatch_table_t &table,
                                       const int engine,
                                       const dpnp::tensor::usm_ndarray &result);
 
+extern sample_fn_ptr_t lookup_sampler(const discrete_dispatch_table_t &table,
+                                      const int engine,
+                                      const bool double_precision);
+
 extern std::pair<sycl::event, sycl::event>
     py_sample(sycl::queue &exec_q,
               const std::vector<std::uint64_t> &seed,
@@ -281,4 +311,34 @@ void def_sampler(py::module_ &m,
           py::arg("params"), py::arg("result"),
           py::arg("depends") = py::list());
 }
+
+// As def_sampler, for a discrete sampler, whose Python function also takes
+// whether to compute in double precision
+template <template <typename> class SamplerT>
+void def_discrete_sampler(py::module_ &m,
+                          const char *name,
+                          const char *doc,
+                          discrete_dispatch_table_t &table)
+{
+    init_discrete_dispatch_table<SamplerT>(table);
+
+    auto sample_pyapi = [&table](sycl::queue &exec_q, const int engine,
+                                 const std::vector<std::uint64_t> &seed,
+                                 const std::uint64_t counter_lo,
+                                 const std::uint64_t counter_hi,
+                                 const py::sequence &params,
+                                 const dpnp::tensor::usm_ndarray &result,
+                                 const bool double_precision,
+                                 const std::vector<sycl::event> &depends) {
+        const sample_fn_ptr_t fn =
+            lookup_sampler(table, engine, double_precision);
+        return py_sample(exec_q, seed, counter_lo, counter_hi, params, result,
+                         depends, fn);
+    };
+
+    m.def(name, sample_pyapi, doc, py::arg("sycl_queue"), py::arg("engine"),
+          py::arg("seed"), py::arg("counter_lo"), py::arg("counter_hi"),
+          py::arg("params"), py::arg("result"), py::arg("double_precision"),
+          py::arg("depends") = py::list());
+}
 } // namespace dpnp::extensions::rng

dpnp/libtensor/extensions/rng/discrete.cpp

@@ -0,0 +1,62 @@
+//*****************************************************************************
+// Copyright (c) 2026, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+// - Neither the name of the copyright holder nor the names of its contributors
+//   may be used to endorse or promote products derived from this software
+//   without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+#include <pybind11/pybind11.h>
+
+#include "common.hpp"
+#include "distributions.hpp"
+
+namespace dpnp::extensions::rng
+{
+static discrete_dispatch_table_t geometric_dispatch_table;
+static discrete_dispatch_table_t logseries_dispatch_table;
+static discrete_dispatch_table_t zipf_dispatch_table;
+static discrete_dispatch_table_t hypergeometric_dispatch_table;
+
+void init_discrete(py::module_ m)
+{
+    def_discrete_sampler<rng_ns::GeometricSampler>(
+        m, "_geometric",
+        "Fill the result with geometric samples of probability p",
+        geometric_dispatch_table);
+
+    def_discrete_sampler<rng_ns::LogseriesSampler>(
+        m, "_logseries",
+        "Fill the result with logarithmic series samples of shape p",
+        logseries_dispatch_table);
+
+    def_discrete_sampler<rng_ns::ZipfSampler>(
+        m, "_zipf", "Fill the result with Zipf samples of exponent a",
+        zipf_dispatch_table);
+
+    def_discrete_sampler<rng_ns::HypergeometricSampler>(
+        m, "_hypergeometric",
+        "Fill the result with hypergeometric samples of nsample draws from "
+        "ngood good and nbad bad items",
+        hypergeometric_dispatch_table);
+}
+} // namespace dpnp::extensions::rng

dpnp/libtensor/extensions/rng/distributions.hpp

@@ -38,4 +38,9 @@ void init_uniform(py::module_ m);
 void init_normal(py::module_ m);
 void init_exponential(py::module_ m);
 void init_gamma(py::module_ m);
+void init_noncentral(py::module_ m);
+void init_integers(py::module_ m);
+void init_poisson(py::module_ m);
+void init_binomial(py::module_ m);
+void init_discrete(py::module_ m);
 } // namespace dpnp::extensions::rng