Basics

This fast-paced section reviews a wide range of standard operations. It assumes the following import declarations

import drjit as dr
from drjit.auto import Float, Array3f, UInt

Creating arrays

Recall that Dr.Jit array types are dynamically sized–for example, Float refers to a 1D array of single precision variables.

The simplest way to create such an array is to call its constructor with a list of explicit values:

a = Float(1, 2, 3, 4)
print(a) # [1, 2, 3, 4]

The constructor also accepts Sequence types (e.g. lists, tuples, NumPy arrays, PyTorch tensors, etc.):

x = Float([1, 2, 3, 4])

Nested array types store several variables—for example, Array3f is just a wrapper around 3 Float instances. They can be passed to the constructor explicitly, or via implicit conversion from constants, lists, etc.

a = Array3f([1, 2], 0, Float(10, 20))
print(a)
# Prints (with 'y' component broadcast to full size)
# [[1, 0, 10],
#  [2, 0, 20]]

Various functions can also create default-initialized arrays:

These always take the desired output type as first argument. You can optionally request a given size along the dynamic axis, e.g.:

b = dr.zeros(Array3f)
print(b.shape) # Prints: (3, 1)

b = dr.zeros(Array3f, shape=(3, 1000))
print(b.shape) # Prints: (3, 1000)

Element access

Use the default array[index] syntax to read/write array entries. Nested static 1-4D arrays further expose equivalent .x / .y / .z / .w members:

a = Array3f(1, 2, 3)
a.x += a.z + a[1]

Static 1-4D arrays also support swizzling, which arbitrarily reorders elements. For example, the following compact compact notation updates and combines entries of a larger array.

a.xy += a.xx + a.zx

Beware that swizzle accesses besides direct assignent or in-place updates create new arrays. As a consequence, the following statement does not modify a as intended, since a.x created a new 1D array.

a.x[a.x < 0] = 0 # Warning: this does not work

Arithmetic operations

Except for a few special cases (e.g., matrix multiplication), arithmetic operations transform arrays element-wise. If needed, the system will implicitly broadcast the operands and promote types.

>>> a = abs(Float(-1.25, 2) + UInt32(1))
>>> type(a)
<class 'drjit.cuda.Float'>
>>> a
[0.25, 3]

In the above example, broadcasting automatically extended the size of the scalar (size-1) array, and the UInt32 type was promoted to Float. Type promotion follows the rules of the C programming language.

Besides built-in arithmetic operators, the following standard functions are available:

The library implements common transcendental functions:

Most of these support real and complex-valued inputs. A subset accepts quaternions (see the section on array types for details). Integer arrays further support the following bit-level operations

Mask operations

Equality and inequality comparisons produce masks (i.e., boolean-valued arrays) with support for binary arithmetic. The dr.select() function blends results from two arrays based on a mask analogous to the ternary (”?”) operator in C/C++.

>>> a = dr.arange(Float, 5) - 3
>>> mask = (a < 0) | (a == 2)
>>> mask
[True, True, False, False, True]
>>> dr.select(mask, -1, a)        # select(mask, true_value, false_value)
[-1, -1, 0, 1, -1]

Masks can also be applied to arrays in order to zero out the False indices by using the & operator.

>>> a = Float([1, 2, 3])
>>> mask = Bool([True, False, True])
>>> a & mask
[1, 0, 3]

Reductions

Reductions use a given operation (e.g., addition) to combine values along one or several dimensions.

By default, they reduce arrays along the leading array dimension. For example, the following reduction is equivalent to a.x + a.y + a.z. By reducing this value once more or specifying axis=None, we can sum over all entries.

>>> a = Array3f([1, 2], [10, 20], [100, 200])
>>> dr.sum(a)
[111, 222]
>>> dr.sum(a, axis=None)
[333]

Accessing memory: gather/scatter

The function dr.gather() fetches values from a 1D array with positions specified by an index array. For example:

>>> buf = Float(10, 20, 30, 40, 50, 60)
>>> index = UInt32(1, 0)
>>> dr.gather(Float, buf, index)
[20, 10]

Note how the operation takes the desired output type as first argument. We can also gather nested arrays (assumed to be flattened in the source 1D array using C-style order) by requesting a different result type.

>>> dr.gather(Array3f, buf, index)
[[40, 50, 60],
 [10, 20, 30]]

Whereas gather reads memory, dr.scatter() realizes the corresponding write operation.

>>> dr.scatter(buf, Array3f(0, 1, 2), UInt32(1))
>>> buf
[10, 20, 30, 0, 1, 2]

Finally, dr.scatter_add() (and the more general dr.scatter_reduce()) atomically accumulates values into an array.

>>> dr.scatter_add(buf, Array3f(100), UInt32(1))
>>> buf
[10, 20, 30, 100, 101, 102]

Random number generation

Dr.Jit was originally developed for Monte Carlo simulation, and programs in that domain require a source of (pseudo-) randomness. For this, the system provides a member of the PCG Family of random number generators by Melissa O’Neill. To try it, import the class drjit.*.PCG32 from your backend of choice and initialize with the desired output array size.

>>> from drjit.auto import PCG32
>>> rng = PCG32(10000)
>>> rng.next_float32()
[0.108379, 0.533909, 0.00684452, .. 9994 skipped .., 0.511698, 0.600626, 0.219648]
>>> rng.next_uint32_bounded(4)
[1, 0, 0, .. 9994 skipped .., 0, 3, 3]

Please see the documentation of this class for a review of its features.