nbxtensor 0.0.8, from the
v0.0.8tag of github.com/NeuroBrix/nbxtensor. Files nameddocs/nbxtensor/…in this text are not part of the nbxtensor repository at this version.
Every example below is executed by the library's own test suite.
from nbxtensor.dtype import itemsize as I
assert I.itemsize("bfloat16") == 2
assert I.itemsize("complex64") == 8
assert I.kind("bfloat16") == "float" and I.kind("uint8") == "uint"
assert I.from_safetensors("BF16") == "bfloat16"
assert I.from_autotune_key("fp16") == "float16"
# a store without float64 holds float32 in its place
assert I.storage_dtype("float64", stores_fp64=False) == "float32"
# a packed int4 weight with float16 scales and minima per group of 128
assert I.representation_bytes("int4-g128-asym", 256) == 128 + 2 * (2 + 2)
try:
I.itemsize("fp8_unknown")
except ValueError as e:
assert "unknown dtype" in str(e)
numpy has no bfloat16: a bfloat16 array on the host is its 16-bit patterns (uint16).
import numpy as np
from nbxtensor.dtype import convert
x = np.array([1.0, 3.14159, -2.5e-3, np.nan], dtype=np.float32)
bits = convert.f32_to_bf16_bits(x) # round to nearest, ties to even
assert bits.dtype == np.uint16 and bits[3] == 0x7FC0 # every NaN becomes the quiet NaN, as torch does
back = convert.bf16_bits_to_f32(bits) # exact widening
assert back[0] == 1.0 and abs(back[1] - 3.14159) < 0.02
half = convert.bf16_bits_to_f16(bits) # bfloat16 -> float16, through float32
assert half.dtype == np.float16
# a caller splitting a large array across workers writes each span into one output
out = np.empty(x.size, dtype=np.float16)
convert.bf16_bits_to_f16(bits[:2], out=out[:2])
convert.bf16_bits_to_f16(bits[2:], out=out[2:])
assert np.array_equal(out[:3], half[:3])
from nbxtensor.dtype.promotion import promoted, float_math_input
assert promoted(["float16", "float32"]) == "float32" # the widest float, past two bytes
assert promoted(["float16", "bfloat16"]) is None # two half dtypes: the caller's precision policy decides
assert float_math_input("int64") == "float32" # rsqrt of an integer reads it as float32
from nbxtensor.dtype.promotion import binary_result
assert binary_result("int8", "uint8") == "int16" # the Array API standard: wide enough for both
assert binary_result("complex64", "float64") == "complex128" # a complex keeps the wider float's precision
assert binary_result("int64", "float16") == "float16" # across kinds (left open by the standard): PyTorch's rule
assert binary_result("bfloat16", "float16") == "float32" # two half formats meet in float32, whatever the order
The dtype an op computes in is decided by a Policy: plain, immutable data handed to the calls that use it (no
context manager, no global switch). Its op classes are PyTorch's autocast tables (nbxtensor/dtype/policy.yml):
some ops compute in float32 whatever the policy, some in its compute dtype, some in the widest of their inputs; an op
in no class keeps its inputs' dtypes.
from nbxtensor.dtype.policy import get_policy
p = get_policy("params=float32,compute=float16,output=float32") # the spelling of JAX's jmp
assert p.compute_dtype("mm", ["float32", "float32"]) == "float16" # a matrix product: the compute dtype
assert p.compute_dtype("softmax", ["float16"]) == "float32" # a softmax: float32, whatever the policy
assert p.compute_dtype("addcmul", ["float16", "float32", "float16"]) == "float32" # the widest input
assert p.fill_value("masked_fill", -3.39e38, "float16") == -65504.0 # a bfloat16 mask sentinel fits float16
assert p.replace(compute="bfloat16").compute == "bfloat16" # a policy never changes: a new one is made
An op's output dtype under a half compute dtype, when it computes in float32 inside (nbxtensor.dtype.outputs):
from nbxtensor.dtype.outputs import amp_fp32_output_dtype, traced_output_is_complex
assert amp_fp32_output_dtype("bfloat16", "bfloat16", safe=False, narrowed=False) == "bfloat16" # follows the graph
assert amp_fp32_output_dtype("float16", "float16", safe=False, narrowed=False) == "float32" # no contract: float32
assert traced_output_is_complex({"output_dtypes": ["complex64"]}) # read from the traced output, never the op's name
Type promotion follows the Array API standard,
and a Python scalar is weak (it takes the array's dtype), as the standard and NumPy's NEP 50 say. Measured by
nbxtensor/tests/test_array_api_promotion.py: the standard's promotion tables, 120 of 120 pairs of the dtypes the
library holds, in both operand orders, and the order-independence of the result over three operands; its scalar rules.
What the standard leaves open, decided once:
| case | nbxtensor | note |
|---|---|---|
| arrays of two kinds (bool, unsigned, signed, float, complex) | the higher kind, its own width | as PyTorch |
| float16 with bfloat16 (the standard defines neither) | float32, whatever the order | as PyTorch |
| a signed integer with uint64 | refused by name | no signed integer holds it |
| a Python float with an integer or bool array | float64 | NEP 50; PyTorch gives float32 |
| a Python complex with an integer array; with a half float array | complex128; complex64 | NEP 50; the table has no complex32 |
| a Python int outside an integer array's range | refused by name | never a silent wrap |
The gap: nbxtensor exposes no array namespace (__array_namespace__) yet, so the standard's own test suite
(array-api-tests) cannot run against it; the conformance above is measured on the standard's tables.
from nbxtensor.dtype.promotion import binary_result, scalar_result
assert binary_result("uint8", "int8") == "int16" # wide enough for both
assert scalar_result("float16", 0.5) == "float16" # a Python scalar is weak
assert scalar_result("int32", 1.5) == "float64" # left open by the standard: NEP 50
from nbxtensor import shape as S
assert S.numel((2, 3, 4)) == 24
assert S.contiguous_strides((2, 3, 4)) == (12, 4, 1)
# the view rule (ATen's computeStride): merging contiguous dims is a view ...
assert S.view_strides((2, 3, 4), (12, 4, 1), (6, 4)) == [4, 1]
# ... flattening a transposed layout is not: None means "copy first"
assert S.view_strides((4, 3), (1, 4), (12,)) is None
# a dim of extent 1 may carry any stride: this decode-step head block is contiguous
assert S.is_contiguous((1, 32, 64, 1), (2048, 64, 1, 64))
assert S.broadcast_shapes((5, 3, 1), (4,)) == (5, 3, 4)
assert S.broadcast_strides((3, 1), (1, 1), (2, 3, 4)) == ((2, 3, 4), (0, 1, 0)) # expand without a copy
assert S.slice_bounds(50, 80, 60) == (50, 60) # clamped, as numpy does
from nbxtensor.dtype.nbx_dtype import NBXDtype, dtype_size, parse_dtype
assert NBXDtype.bfloat16.value == 1 # values are frozen: new members are appended
assert dtype_size(NBXDtype.complex64) == 8 # read from the one table, never written twice
assert parse_dtype("torch.float16") is NBXDtype.float16 # a torch spelling, read as text
assert parse_dtype("half") is NBXDtype.float16 # an alias the table carries
The library measures the facts of the machine it runs on and has a default for every setting, so import nbxtensor
works with no application around it:
nbxtensor/backend/devices/<vendor>/<name>.yml, read by nbxtensor.backend.device_files);nbxtensor.backend.host_memory.memory_state), whether the devicenbxtensor.backend.device_facts);nbxtensor.ops.tiling_policy.set_tiling, default: the library's ops/tiling.yml), the directory of proven kernelnbxtensor.ops.autotune_settings.set_store, a path or a callable returning one, default: your cache, ~/.cache/nbxtensor/certified/,nbxtensor.memory.set_limit, default:None restores a default.from pathlib import Path
from nbxtensor.backend.device_names import arch_string
from nbxtensor.backend.host_memory import memory_state
from nbxtensor.ops import autotune_settings, tiling_policy
from nbxtensor.backend import device_files
m4 = device_files.read("apple", "apple_m4_pro") # one device's description, from the package itself
assert m4["compute_capability"] == "apple-m4-pro" and m4["memory"]["max_shared_memory_per_block"] == 32768
assert arch_string("Apple M4 Pro") == "apple-m4-pro" # a device's key: the library's own rule
assert memory_state().source # the host reading names where it came from
autotune_settings.set_store("/tmp/my-proven-settings") # your own store of proven kernel settings
assert autotune_settings.directory() == Path("/tmp/my-proven-settings")
autotune_settings.set_store(None) # back to the default in your cache
table = tiling_policy.tiling_table() # the split sizes in force (a copy)
tiling_policy.set_tiling(table) # a whole table, as tiling.yml writes it
tiling_policy.set_tiling(None) # back to the library's file
A memory limit is a setting, not a service: nbxtensor.memory.set_limit holds the process to a card of that many
bytes (the allocator refuses what such a card refuses, and its pool parks a fraction of the limit); absent, the
physical card is the limit.
from nbxtensor import memory
memory.set_limit(8 << 30) # behave as a card of 8 GiB on every device
memory.set_limit(None) # the physical card again
Some services are still passed through nbxtensor.host while the boundary is cleaned before 0.0.1 (the device
description, the operand decisions of the matrix unit and attention, and where a backend build may live).
docs/nbxtensor/API.md section 11 lists each one, its fate (generic in the library, or back in the application) and
its state.
An application may also watch every kernel launch without running it — a dry run, to list the kernel keys a model
forms without a device: nbxtensor.ops.launch_observer.set_launch_observer installs an observer (active(),
record(tuned, key), shadow_run(...)); with none installed, every launch runs.
from nbxtensor.ops import launch_observer
assert launch_observer.launch_observer().active() is False # no observer: nothing runs dry
Host work over a large array runs on a thread pool, one contiguous piece per thread, each piece a numpy call (numpy
releases the GIL in its loops). The pool's size is measured on the machine (Apple silicon: its performance cores) and
can be set; under a threshold measured for each kind of work, the work runs inline. The BLAS numpy calls is held to
the same cores: as many threads as the pool outside it, one per piece inside parallel_for, so the two never
oversubscribe the machine (nbxtensor.backend.cpu.blas: Apple's Accelerate per thread, OpenBLAS for the process).
import numpy as np
from nbxtensor.backend.cpu import pool
assert pool.threads() == pool.measured_threads() >= 1 # this machine's compute cores, measured
x = np.linspace(-3, 3, 1 << 20, dtype=np.float32)
out = np.empty_like(x)
probe = lambda n: (lambda a, b: np.exp(x[a:b], out=out[a:b])) # how the threshold for "exp" is measured
pool.parallel_for(x.size, lambda a, b: np.exp(x[a:b], out=out[a:b]), "exp", probe)
assert out.tobytes() == np.exp(x).tobytes() # the same bytes as one numpy call
pool.set_threads(2) # an application's own size
assert pool.threads() == 2
pool.set_threads(None) # back to the measurement
The host conversions of a weight at load are split the same way, across the loader's workers
(nbxtensor.io.safetensors.to_bf16_bits: float32 or float16 rounded to bfloat16, the bytes of one call).
nbxtensor speaks DLPack and nothing else: Tensor.__dlpack__ / __dlpack_device__ and nbxtensor.from_dlpack(x)
for any producer. The other library imports itself (torch.from_dlpack(t)); nbxtensor never imports torch. Planned
with the tensor class (API.md §4).