nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/rand_op.py

@@ -0,0 +1,52 @@
+"""Random number generation — pure @triton.jit kernels.
+
+rand_kernel: uniform [0, 1) using tl.rand(seed, offset)
+randn_kernel: standard normal N(0,1) using tl.randn(seed, offset)
+
+Both use Triton's built-in PRNG which implements Philox counter-based RNG.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def rand_kernel(
+    output_ptr,
+    n_elements,
+    seed,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Fill output with uniform random values in [0, 1).
+
+    Args:
+        output_ptr: pointer to output tensor (flattened)
+        n_elements: total number of elements
+        seed: RNG seed (int)
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    r = tl.rand(seed, offset)
+    tl.store(output_ptr + offset, r, mask=mask)
+
+@triton.jit
+def randn_kernel(
+    output_ptr,
+    n_elements,
+    seed,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Fill output with standard normal random values N(0, 1).
+
+    Args:
+        output_ptr: pointer to output tensor (flattened)
+        n_elements: total number of elements
+        seed: RNG seed (int)
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    r = tl.randn(seed, offset)
+    tl.store(output_ptr + offset, r, mask=mask)

nbxtensor/ops/random_wrappers.py

@@ -0,0 +1,366 @@
+"""Random wrappers: rand, randn, uniform, multinomial, dropout, the seeded draws. Cut out of wrappers.py by family (M2, 2026-10-29), unchanged; shared launch layer and run context stay in wrappers.py."""
+
+from __future__ import annotations
+
+import os
+import triton
+from nbxtensor.ops.autotune_bucket import bucket_of as _bucket_of
+from nbxtensor.ops import launch_keys as _lk
+from nbxtensor.tensor import NBXTensor
+from nbxtensor.dtype.nbx_dtype import NBXDtype, dtype_size
+from nbxtensor.memory.allocator import DeviceAllocator
+from nbxtensor.backend.runtime import set_device
+from nbxtensor.shape import MAX_NDIM as _NBX_MAX_NDIM, broadcast_shapes, max_linear_offset
+from nbxtensor.dtype.itemsize import canonical_name
+from nbxtensor.dtype.promotion import binary_result, float_math_input
+from nbxtensor.dtype.nbx_dtype import parse_dtype
+from nbxtensor.memory.allocator import DeviceOOMError
+from nbxtensor.backend.capabilities import fa_min_tile as _fa_min_tile
+from nbxtensor.backend.capabilities import reduce_tile_max as _reduce_tile_max
+from nbxtensor.backend.faults import device_fault_buffer, device_fault_code_cached, fault_channel
+from nbxtensor.tensor import _upload_int64_array as _nbx_upload_int64, _saturating_cast as _nbx_saturating_cast
+from nbxtensor.ops._configs import sdpa_block_ceiling as _sdpa_block_ceiling
+from nbxtensor.ops._configs import sdpa_launch_meta as _sdpa_launch_meta
+from nbxtensor.ops._configs import matrix_unit as _matrix_unit, matrix_unit_mm_tile as _matrix_unit_mm_tile
+from nbxtensor.ops._configs import largest_tile_within_smem as _largest_tile_within_smem
+from nbxtensor.backend import launcher as _nbx_launcher
+from nbxtensor.backend.launcher import install as _install_launcher
+from nbxtensor.ops.relu import relu_forward_kernel
+from nbxtensor.ops.silu import silu_forward_kernel
+from nbxtensor.ops.gelu import gelu_forward_kernel
+from nbxtensor.ops.sigmoid import sigmoid_forward_kernel
+from nbxtensor.ops.tanh import tanh_forward_kernel
+from nbxtensor.ops.hardsigmoid import hardsigmoid_forward_kernel
+from nbxtensor.ops.hardswish import hardswish_forward_kernel
+from nbxtensor.ops.leaky_relu import leaky_relu_forward_kernel
+from nbxtensor.ops.elu import elu_forward_kernel
+from nbxtensor.ops.mish import mish_forward_kernel
+from nbxtensor.ops.selu import selu_forward_kernel
+from nbxtensor.ops.neg import neg_forward_kernel
+from nbxtensor.ops.exp import exp_forward_kernel
+from nbxtensor.ops.sin import sin_forward_kernel
+from nbxtensor.ops.cos import cos_forward_kernel
+from nbxtensor.ops.rsqrt import rsqrt_forward_kernel
+from nbxtensor.ops.sqrt import sqrt_forward_kernel
+from nbxtensor.ops.abs import abs_forward_kernel
+from nbxtensor.ops.log import log_forward_kernel
+from nbxtensor.ops.reciprocal import reciprocal_forward_kernel
+from nbxtensor.ops.pow import pow_forward_kernel
+from nbxtensor.ops.clamp import clamp_forward_kernel
+from nbxtensor.ops.erf import erf_forward_kernel
+from nbxtensor.ops.copy import copy_forward_kernel
+from nbxtensor.ops.add import add_forward_kernel, add_scalar_kernel, add_scalar_dev_kernel, add_bias_broadcast_kernel
+from nbxtensor.ops.binary_strided import add_strided_nd_kernel, mul_strided_nd_kernel, sub_strided_nd_kernel, div_strided_nd_kernel, add_scalar_strided_nd_kernel, mul_scalar_strided_nd_kernel
+from nbxtensor.ops.mul import mul_forward_kernel, mul_scalar_kernel, mul_scalar_dev_kernel
+from nbxtensor.ops.div import div_forward_kernel, div_scalar_kernel, div_scalar_dev_kernel, int_div_kernel
+from nbxtensor.ops.sub import sub_forward_kernel, rsub_forward_kernel
+from nbxtensor.ops.where import where_forward_kernel
+from nbxtensor.ops.maximum import maximum_forward_kernel
+from nbxtensor.ops.minimum import minimum_forward_kernel
+from nbxtensor.ops.masked_fill import masked_fill_forward_kernel
+from nbxtensor.ops.gt import gt_forward_kernel, gt_scalar_kernel
+from nbxtensor.ops.ge import ge_forward_kernel, ge_scalar_kernel
+from nbxtensor.ops.lt import lt_forward_kernel, lt_scalar_kernel
+from nbxtensor.ops.eq import eq_forward_kernel, eq_scalar_kernel
+from nbxtensor.ops.ne import ne_forward_kernel, ne_scalar_kernel
+from nbxtensor.ops.layernorm import layer_norm_forward_kernel
+from nbxtensor.ops.rmsnorm import rms_norm_forward_kernel
+from nbxtensor.ops.softmax import softmax_forward_kernel, softmax_forward_looped_kernel
+from nbxtensor.ops.matmul import matmul_kernel, addmm_kernel
+from nbxtensor.ops._autotune_policy import NBX_FORCE_FP32_ACCUM
+from nbxtensor.ops.fused_moe import fused_moe_kernel, silu_and_mul_kernel, silu_mul_split_kernel
+from nbxtensor.ops.dtype_convert import bf16_to_fp16_kernel
+from nbxtensor.ops.embedding import embedding_kernel, EMBEDDING_OOB
+from nbxtensor.ops.sum import sum_forward_kernel
+from nbxtensor.ops.mean import mean_forward_kernel
+from nbxtensor.ops.amax import amax_forward_kernel
+from nbxtensor.ops.glu import glu_forward_kernel
+from nbxtensor.ops.le import le_forward_kernel, le_scalar_kernel
+from nbxtensor.ops.softplus import softplus_forward_kernel
+from nbxtensor.ops.dropout import dropout_inference_kernel
+from nbxtensor.ops.upsample_nearest2d import upsample_nearest2d_kernel
+from nbxtensor.ops.groupnorm import group_norm_forward_kernel
+from nbxtensor.ops.index_select import index_select_kernel, INDEX_SELECT_OOB
+from nbxtensor.ops.triu import triu_kernel, triu_batch_kernel
+from nbxtensor.ops.tril import tril_kernel, tril_batch_kernel
+from nbxtensor.ops.argmax import argmax_kernel_1, argmax_kernel_2, argmax_kernel_inner
+from nbxtensor.ops.argmin import argmin_kernel_1, argmin_kernel_2, argmin_kernel
+from nbxtensor.ops.conv2d import conv2d_forward_kernel
+from nbxtensor.ops.depthwise_conv2d import depthwise_conv2d_kernel
+from nbxtensor.ops.batch_norm import batch_norm_forward_kernel
+from nbxtensor.ops.cumsum import scan_part_sum_kernel, add_base_sum_kernel
+from nbxtensor.ops.scatter_op import scatter_kernel, scatter_add_kernel, scatter_reduce_amax_kernel, scatter_reduce_amin_kernel
+from nbxtensor.ops.gather_op import gather_kernel
+from nbxtensor.ops.avg_pool2d import avg_pool2d_forward_kernel
+from nbxtensor.ops.max_pool2d import max_pool2d_forward_kernel
+from nbxtensor.ops.max_pool3d import max_pool3d_forward_kernel
+from nbxtensor.ops.cross_entropy import cross_entropy_loss_forward_kernel
+from nbxtensor.ops.exp2 import exp2_forward_kernel
+from nbxtensor.ops.tan import tan_forward_kernel
+from nbxtensor.ops.celu import celu_forward_kernel
+from nbxtensor.ops.log_sigmoid import log_sigmoid_forward_kernel
+from nbxtensor.ops.isfinite import isfinite_forward_kernel
+from nbxtensor.ops.isinf import isinf_forward_kernel
+from nbxtensor.ops.isnan import isnan_forward_kernel
+from nbxtensor.ops.nan_to_num import nan_to_num_forward_kernel
+from nbxtensor.ops.threshold import threshold_forward_kernel
+from nbxtensor.ops.floor import floor_forward_kernel, ceil_forward_kernel, round_forward_kernel, trunc_forward_kernel
+from nbxtensor.ops.addcdiv import addcdiv_forward_kernel
+from nbxtensor.ops.addcmul import addcmul_forward_kernel
+from nbxtensor.ops.all_reduce import all_kernel_mid, all_kernel_result, all_kernel_dim
+from nbxtensor.ops.any_reduce import any_kernel_mid, any_kernel_result, any_kernel_dim
+from nbxtensor.ops.bitwise_and import bitwise_and_forward_kernel
+from nbxtensor.ops.bitwise_or import bitwise_or_forward_kernel
+from nbxtensor.ops.bitwise_not import bitwise_not_forward_kernel
+from nbxtensor.ops.logical_and import logical_and_forward_kernel
+from nbxtensor.ops.logical_or import logical_or_forward_kernel
+from nbxtensor.ops.logical_not import logical_not_forward_kernel
+from nbxtensor.ops.lerp import lerp_tensor_forward_kernel, lerp_scalar_head_kernel, lerp_scalar_tail_kernel
+from nbxtensor.ops.dot_op import dot_kernel_small, dot_kernel_partial, dot_kernel_reduce
+from nbxtensor.ops.flip_op import flip_1d_kernel, flip_strided_kernel
+from nbxtensor.ops.cat_op import cat_copy_kernel_4
+from nbxtensor.ops.fill_op import fill_kernel
+from nbxtensor.ops.arange_op import arange_kernel
+from nbxtensor.ops.prod import prod_kernel_mid, prod_kernel_result, prod_kernel
+from nbxtensor.ops.min_reduce import min_kernel_mid, min_kernel_result, min_kernel
+from nbxtensor.ops.max_reduce import max_kernel_mid, max_kernel_result, max_kernel
+from nbxtensor.ops.baddbmm_op import baddbmm_kernel
+from nbxtensor.ops.addmv_op import addmv_kernel
+from nbxtensor.ops.mse_loss import mse_loss_partial_kernel, mse_loss_reduce_kernel, mse_loss_none_kernel
+from nbxtensor.ops.nllloss import nll_loss_forward_kernel
+from nbxtensor.ops.std import std_map_kernel, std_reduce_kernel, std_dim_kernel
+from nbxtensor.ops.var import var_kernel_1, var_kernel_2, var_welford_kernel
+from nbxtensor.ops.index_add import index_add_gather_kernel
+from nbxtensor.ops.index_put_op import index_put_kernel, INDEX_PUT_OOB
+from nbxtensor.ops.sort_op import (radix_sort_histogram_kernel,
+                          radix_sort_scatter_kernel,
+                          radix_sort_tile_counts_kernel,
+                          radix_sort_tile_prefix_kernel)
+from nbxtensor.ops.rope import rope_forward_kernel
+from nbxtensor.ops.pixel_shuffle import (
+    pixel_shuffle_kernel,
+    pixel_shuffle_broadcast_aware_kernel,
+)
+from nbxtensor.ops.pixel_unshuffle import pixel_unshuffle_kernel
+from nbxtensor.ops.upsample_bilinear2d import upsample_bilinear2d_kernel
+from nbxtensor.ops.adaptive_avg_pool2d import adaptive_avg_pool2d_kernel
+from nbxtensor.ops.conv_depthwise2d import conv_depthwise2d_kernel
+from nbxtensor.ops.scaled_softmax import scaled_softmax_kernel
+from nbxtensor.ops.vector_norm import (
+    l2_norm_kernel, l1_norm_kernel, linf_norm_kernel, l0_norm_kernel, lp_norm_kernel,
+    l2_norm_pass1_kernel, l2_norm_pass2_kernel,
+    linf_norm_pass1_kernel, linf_norm_pass2_kernel,
+)
+from nbxtensor.ops.addr import addr_kernel
+from nbxtensor.ops.logical_xor import logical_xor_kernel
+from nbxtensor.ops.bitwise_left_shift import bitwise_left_shift_kernel
+from nbxtensor.ops.bitwise_right_shift import bitwise_right_shift_kernel
+from nbxtensor.ops.rand_op import rand_kernel, randn_kernel
+from nbxtensor.ops.clamp_min import clamp_min_forward_kernel
+from nbxtensor.ops.conv_transpose2d import conv_transpose2d_kernel
+from nbxtensor.ops.cummax import scan_part_max_kernel
+from nbxtensor.ops.cummin import scan_part_min_kernel
+from nbxtensor.ops.count_nonzero import count_nonzero_kernel
+from nbxtensor.ops.trace_op import trace_kernel
+from nbxtensor.ops.linspace_op import linspace_kernel
+from nbxtensor.ops.full_op import full_kernel
+from nbxtensor.ops.stack_op import stack_copy_kernel
+from nbxtensor.ops.repeat_op import repeat_1d_kernel
+from nbxtensor.ops.upsample_nearest3d import upsample_nearest3d_kernel
+from nbxtensor.ops.weight_norm import weight_norm_kernel_first, weight_norm_kernel_last
+from nbxtensor.ops.repeat_interleave import repeat_interleave_tensor_kernel
+from nbxtensor.ops.elementwise_launch import ew_meta as _ew_meta, ew_grid as _ew_grid
+from nbxtensor.ops import tiling_sizes as _ts
+import math
+from nbxtensor.ops import elementwise_wrappers as _ew
+from nbxtensor.ops import reduction_wrappers as _red
+from nbxtensor.ops import wrappers as _core
+
+
+def dropout_wrapper(x, p: float = 0.5, train: bool = False) :
+    """Dropout at inference = identity. Training dropout not implemented in Triton yet."""
+    if not train or p == 0.0:
+        # Inference mode: passthrough
+        x = x.contiguous()
+        output = NBXTensor.empty_like(x)
+        set_device(x)
+        dropout_inference_kernel[_ew._1d_grid(x.numel())](x, output, x.numel(), **_ew_meta(x.numel()))
+        return output
+    raise NotImplementedError("Training dropout with Triton RNG not implemented yet")
+
+
+# ===========================================================================
+# RNG OPS — Random number generation via Triton kernels
+# ===========================================================================
+
+def _draw_seed() -> int:
+    """Per-draw Philox kernel seed for the random wrappers.
+
+    Run-scoped stream when armed (reproducible, --seed-driven; the Triton
+    mirror of the compiled engine's run-scoped generator — see
+    nbxtensor/random/rng_stream.py), else the historical unseeded fallback. The
+    NBX_FORCE_RAND_SEED pin (rng_pin) is checked by the callers BEFORE any
+    kernel draw and keeps precedence over this stream.
+    """
+    from nbxtensor.random import rng_stream
+    if rng_stream.active():
+        return rng_stream.next_seed()
+    return __import__("random").randint(0, 2**31 - 1)
+
+
+def rand_wrapper(*args, **kwargs) :
+    """rand(size, ...) → uniform [0, 1) via Triton kernel."""
+    from nbxtensor.ops.rand_op import rand_kernel
+    from nbxtensor.random.rng_pin import pinned_seed, pinned_uniform
+    # ATen signature: rand(SymInt[] size, ...)
+    size = args[0] if args else kwargs.get('size', [])
+    if isinstance(size, NBXTensor):
+        size = size.tolist()
+    dtype = kwargs.get('dtype', NBXDtype.float32)
+    device = kwargs.get('device', None)
+    if device is None:
+        device = str('cuda')
+    if pinned_seed() is not None:
+        return NBXTensor.from_numpy(pinned_uniform(size)).to(dtype)
+    drawn = _ew._stream_draw("uniform", size, dtype, device)
+    if drawn is not None:
+        return drawn
+    output = NBXTensor.empty(size, dtype=dtype, device=device)
+    n = output.numel()
+    if n == 0:
+        return output
+    seed = _draw_seed()
+    set_device(output)
+    rand_kernel[_ew._1d_grid(n)](output, n, seed, **_ew_meta(n))
+    return output
+
+
+def randn_wrapper(*args, **kwargs) :
+    """randn(size, ...) → standard normal N(0,1) via Triton kernel."""
+    from nbxtensor.ops.rand_op import randn_kernel
+    size = args[0] if args else kwargs.get('size', [])
+    if isinstance(size, NBXTensor):
+        size = size.tolist()
+    dtype = kwargs.get('dtype', NBXDtype.float32)
+    device = kwargs.get('device', None)
+    if device is None:
+        device = str('cuda')
+    drawn = _ew._stream_draw("normal", size, dtype, device)
+    if drawn is not None:
+        return drawn
+    output = NBXTensor.empty(size, dtype=dtype, device=device)
+    n = output.numel()
+    if n == 0:
+        return output
+    seed = _draw_seed()
+    set_device(output)
+    randn_kernel[_ew._1d_grid(n)](output, n, seed, **_ew_meta(n))
+    return output
+
+
+def rand_like_wrapper(x, **kwargs) :
+    """rand_like(tensor) → uniform [0, 1) same shape/dtype/device."""
+    from nbxtensor.ops.rand_op import rand_kernel
+    dtype = kwargs.get('dtype', x.dtype)
+    device = kwargs.get('device', x.device)
+    drawn = _ew._stream_draw("uniform", x.shape, kwargs.get('dtype', x._dtype), device)
+    if drawn is not None:
+        return drawn
+    output = NBXTensor.empty(x.shape, dtype=dtype, device=device)
+    n = output.numel()
+    if n == 0:
+        return output
+    seed = _draw_seed()
+    set_device(output)
+    rand_kernel[_ew._1d_grid(n)](output, n, seed, **_ew_meta(n))
+    return output
+
+
+def randn_like_wrapper(x, **kwargs) :
+    """randn_like(tensor) → standard normal N(0,1) same shape/dtype/device."""
+    from nbxtensor.ops.rand_op import randn_kernel
+    from nbxtensor.random.rng_pin import pinned_seed, pinned_normal
+    dtype = kwargs.get('dtype', x.dtype)
+    device = kwargs.get('device', x.device)
+    if pinned_seed() is not None:
+        return NBXTensor.from_numpy(pinned_normal(x.shape)).to(dtype)
+    drawn = _ew._stream_draw("normal", x.shape, kwargs.get('dtype', x._dtype), device)
+    if drawn is not None:
+        return drawn
+    output = NBXTensor.empty(x.shape, dtype=dtype, device=device)
+    n = output.numel()
+    if n == 0:
+        return output
+    seed = _draw_seed()
+    set_device(output)
+    randn_kernel[_ew._1d_grid(n)](output, n, seed, **_ew_meta(n))
+    return output
+
+
+def uniform_wrapper(x, low: float = 0.0, high: float = 1.0,
+                    **kwargs) :
+    """uniform_(tensor, low, high) → fill with uniform [low, high)."""
+    from nbxtensor.ops.rand_op import rand_kernel
+    x = _core._ensure_cuda(x).contiguous()
+    n = x.numel()
+    if n == 0:
+        return x
+    seed = _draw_seed()
+    rand_kernel[_ew._1d_grid(n)](x, n, seed, **_ew_meta(n))
+    # Transform [0,1) → [low, high)
+    if high - low != 1.0:
+        x.mul_(high - low)
+    if low != 0.0:
+        x.add_(low)
+    return x
+
+
+def multinomial_wrapper(x, num_samples: int,
+                        replacement: bool = False, **kwargs) :
+    """multinomial(probs, num_samples) → sampled indices.
+
+    Uses Gumbel-max trick: argmax(log(probs) + Gumbel) = categorical sample.
+    Gumbel noise = -log(-log(uniform)).
+    """
+    from nbxtensor.ops.rand_op import rand_kernel
+    x = _core._ensure_cuda(x)
+    # Normalize probabilities
+    probs = x.float()
+    if probs.ndim == 1:
+        probs = probs.unsqueeze(0)
+        squeeze_out = True
+    else:
+        squeeze_out = False
+
+    batch_size, n_categories = probs.shape
+
+    if num_samples == 1:
+        # Single sample: Gumbel-max trick
+        # gumbel = -log(-log(uniform(0,1)))
+        gumbel = NBXTensor.empty_like(probs)
+        n = gumbel.numel()
+        seed = _draw_seed()
+        rand_kernel[_ew._1d_grid(n)](gumbel, n, seed, **_ew_meta(n))
+        gumbel = _ew.clamp(gumbel, min_val=1e-10)
+        gumbel = _ew.neg(_ew.log_wrapper(_ew.neg(_ew.log_wrapper(gumbel))))
+        log_probs = _ew.clamp(_ew.log_wrapper(probs), min_val=-1e30)
+        scores = _ew.add(log_probs, gumbel)
+        indices = _red.argmax_wrapper(scores, dim=-1, keepdim=True)
+    else:
+        # Multiple samples: repeated Gumbel-max
+        all_indices = []
+        for _ in range(num_samples):
+            gumbel = NBXTensor.empty_like(probs)
+            n = gumbel.numel()
+            seed = _draw_seed()
+            rand_kernel[_ew._1d_grid(n)](gumbel, n, seed, **_ew_meta(n))
+            gumbel = _ew.clamp(gumbel, min_val=1e-10)
+            gumbel = _ew.neg(_ew.log_wrapper(_ew.neg(_ew.log_wrapper(gumbel))))
+            log_probs = _ew.clamp(_ew.log_wrapper(probs), min_val=-1e30)
+            scores = _ew.add(log_probs, gumbel)
+            all_indices.append(_red.argmax_wrapper(scores, dim=-1, keepdim=True))
+        indices = NBXTensor.cat(all_indices, dim=-1)
+
+    if squeeze_out:
+        indices = indices.squeeze(0)
+    return indices.to(NBXDtype.int64)

nbxtensor/ops/reciprocal.py

@@ -0,0 +1,17 @@
+"""Reciprocal (1/x) — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def reciprocal_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, 1.0 / x.to(tl.float32), mask=mask)