Reference#

Array creation#

drjit.zeros(dtype, shape=1)#

Return a zero-initialized instance of the desired type and shape

This function can create zero-initialized instances of various types. In particular, dtype can be:

  • A Dr.Jit array type like drjit.cuda.Array2f. When shape specifies a sequence, it must be compatible with static dimensions of the dtype. For example, dr.zeros(dr.cuda.Array2f, shape=(3, 100)) fails, since the leading dimension is incompatible with drjit.cuda.Array2f. When shape is an integer, it specifies the size of the last (dynamic) dimension, if available.

  • A tensorial type like drjit.scalar.TensorXf. When shape specifies a sequence (list/tuple/..), it determines the tensor rank and shape. When shape is an integer, the function creates a rank-1 tensor of the specified size.

  • A custom data structure. In this case, drjit.zeros() will invoke itself recursively to zero-initialize each field of the data structure.

  • A scalar Python type like int, float, or bool. The shape parameter is ignored in this case.

Note that when dtype refers to a scalar mask or a mask array, it will be initialized to False as opposed to zero.

Parameters:
  • dtype (type) – Desired Dr.Jit array type, Python scalar type, or custom data structure.

  • shape (Sequence[int] | int) – Shape of the desired array

Returns:

A zero-initialized instance of type dtype.

Return type:

object

drjit.ones(dtype, shape=1)#

Return a constant-valued instance of the desired type and shape filled with ones.

This function can create constant-valued instances of various types. In particular, dtype can be:

specifies a sequence, it must be compatible with static dimensions of the dtype. For example, dr.ones(dr.cuda.Array2f, shape=(3, 100)) fails, since the leading dimension is incompatible with

drjit.cuda.Array2f. When shape is an integer, it specifies

the size of the last (dynamic) dimension, if available.

  • A tensorial type like drjit.scalar.TensorXf. When shape

specifies a sequence (list/tuple/..), it determines the tensor rank and shape. When shape is an integer, the function creates a rank-1 tensor of the specified size.

  • A custom data structure. In this case,

drjit.ones() will invoke itself recursively to initialize each field of the data structure.

  • A scalar Python type like int, float, or bool. The shape

parameter is ignored in this case.

Parameters:
  • dtype (type) – Desired Dr.Jit array type, Python scalar type, or custom data structure.

  • shape (Sequence[int] | int) – Shape of the desired array

Returns:

A instance of type dtype filled with ones

Return type:

object

drjit.full(dtype, value, shape=1)#

Return a constant-valued instance of the desired type and shape

This function can create constant-valued instances of various types. In particular, dtype can be:

specifies a sequence, it must be compatible with static dimensions of the dtype. For example, dr.full(dr.cuda.Array2f, value=1.0, shape=(3, 100)) fails, since the leading dimension is incompatible with drjit.cuda.Array2f. When shape is an integer, it specifies the size of the last (dynamic) dimension, if available.

  • A tensorial type like drjit.scalar.TensorXf. When shape

specifies a sequence (list/tuple/..), it determines the tensor rank and shape. When shape is an integer, the function creates a rank-1 tensor of the specified size.

  • A custom data structure. In this case,

drjit.full() will invoke itself recursively to initialize each field of the data structure.

  • A scalar Python type like int, float, or bool. The shape

parameter is ignored in this case.

Parameters:
  • dtype (type) – Desired Dr.Jit array type, Python scalar type, or custom data structure.

  • value (object) – An instance of the underlying scalar type (float/int/bool, etc.) that will be used to initialize the array contents.

  • shape (Sequence[int] | int) – Shape of the desired array

Returns:

A instance of type dtype filled with value

Return type:

object

drjit.identity(dtype, size=1)#

Return the identity array of the desired type and size

This function can create identity instances of various types. In particular, dtype can be:

  • A Dr.Jit matrix type (like drjit.cuda.Matrix4f).

  • A Dr.Jit complex type (like drjit.cuda.Quaternion4f).

  • Any other Dr.Jit array type. In this case this function is equivalent to full(dtype, 1, size)

  • A scalar Python type like int, float, or bool. The size parameter is ignored in this case.

Parameters:
  • dtype (type) – Desired Dr.Jit array type, Python scalar type, or custom data structure.

  • value (object) – An instance of the underlying scalar type (float/int/bool, etc.) that will be used to initialize the array contents.

  • size (int) – Size of the desired array | matrix

Returns:

The identity array of type dtype of size size

Return type:

object

drjit.arange(dtype, start=None, stop=None, step=1)#

This function generates an integer sequence on the interval [start, stop) with step size step, where start = 0 and step = 1 if not specified.

Parameters:
  • dtype (type) – Desired Dr.Jit array type. The dtype must refer to a dynamically sized 1D Dr.Jit array such as drjit.scalar.ArrayXu or drjit.cuda.Float.

  • start (int) – Start of the interval. The default value is 0.

  • stop/size (int) – End of the interval (not included). The name of this parameter differs between the two provided overloads.

  • step (int) – Spacing between values. The default value is 1.

Returns:

The computed sequence of type dtype.

Return type:

object

drjit.linspace(dtype, start, stop, num=1, endpoint=True)#

This function generates an evenly spaced floating point sequence of size num covering the interval [start, stop].

Parameters:
  • dtype (type) – Desired Dr.Jit array type. The dtype must refer to a dynamically sized 1D Dr.Jit floating point array, such as drjit.scalar.ArrayXf or drjit.cuda.Float.

  • start (float) – Start of the interval.

  • stop (float) – End of the interval.

  • num (int) – Number of samples to generate.

  • endpoint (bool) – Should the interval endpoint be included? The default is True.

Returns:

The computed sequence of type dtype.

Return type:

object

drjit.tile(arg, count: int)#

This function constructs an Dr.Jit array by repeating arg count times.

Parameters:
  • arg (drjit.ArrayBase) – A Dr.Jit type

  • count (int) – Number of repetitions

Returns:

The tiled output array.

Return type:

object

drjit.repeat(array, count: int)#

This function constructs an Dr.Jit array by repeating the elements of arg count times.

Parameters:
  • arg (drjit.ArrayBase) – A Dr.Jit type

  • count (int) – Number of repetitions for the elements

Returns:

Output array where the elements where repeated.

Return type:

object

drjit.meshgrid(*args, indexing='xy')#

Creates a grid coordinates based on the coordinates contained in the provided one-dimensional arrays.

The indexing keyword argument allows this function to support both matrix and Cartesian indexing conventions. If given the string ‘ij’, it will return a grid coordinates with matrix indexing. If given ‘xy’, it will return a grid coordinates with Cartesian indexing.

import drjit as dr

x, y = dr.meshgrid(
    dr.arange(dr.llvm.UInt, 4),
    dr.arange(dr.llvm.UInt, 4)
)

# x = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
# y = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]
Parameters:
  • args (drjit.ArrayBase) – Dr.Jit one-dimensional coordinate arrays

  • indexing (str) – Specifies the indexing conventions

Returns:

Grid coordinates

Return type:

tuple

Horizontal operations#

These operations are horizontal in the sense that [..]

drjit.gather(dtype, source, index, active=True)#

Gather values from a flat array or nested data structure

This function performs a gather (i.e., indirect memory read) from source at position index. It expects a dtype argument and will return an instance of this type. The optional active argument can be used to disable some of the components, which is useful when not all indices are valid; the corresponding output will be zero in this case.

This operation can be used in the following different ways:

1. When dtype is a 1D Dr.Jit array like drjit.llvm.ad.Float, this operation implements a parallelized version of the Python array indexing expression source[index] with optional masking. Example:

source = dr.cuda.Float([...])
index = dr.cuda.UInt([...]) # Note: negative indices are not permitted
result = dr.gather(dtype=type(source), source=source, index=index)
  1. When dtype is a more complex type (e.g. a custom source structure,

    nested Dr.Jit array, tuple, list, dictionary, etc.), the behavior depends:

  • When type(source) matches dtype, the the gather operation threads

    through entries and invokes itself recursively. For example, the gather operation in

    result = dr.cuda.Array3f(...)
    index = dr.cuda.UInt([...])
    result = dr.gather(dr.cuda.Array3f, source, index)
    

    is equivalent to

    result = dr.cuda.Array3f(
        dr.gather(dr.cuda.Float, source.x, index),
        dr.gather(dr.cuda.Float, source.y, index),
        dr.gather(dr.cuda.Float, source.z, index)
    )
    
  • Otherwise, the operation reconstructs the requested dtype from a flat

    source array, using C-style ordering with a suitably modified index. For example, the gather below reads 3D vectors from a 1D array.

    source = dr.cuda.Float([...])
    index = dr.cuda.UInt([...])
    result = dr.gather(dr.cuda.Array3f, source, index)
    

    and is equivalent to

    result = dr.cuda.Vector3f(
        dr.gather(dr.cuda.Float, source, index*3 + 0),
        dr.gather(dr.cuda.Float, source, index*3 + 1),
        dr.gather(dr.cuda.Float, source, index*3 + 2)
    )
    

Danger

The indices provided to this operation are unchecked. Out-of-bounds reads are undefined behavior (if not disabled via the active parameter) and may crash the application. Negative indices are not permitted.

Parameters:
  • dtype (type) – The desired output type (typically equal to type(source), but other variations are possible as well, see the description above.)

  • source (object) – The object from which data should be read (typically a 1D Dr.Jit array, but other variations are possible as well, see the description above.)

  • index (object) – a 1D dynamic unsigned 32-bit Dr.Jit array (e.g., drjit.scalar.ArrayXu or drjit.cuda.UInt) specifying gather indices. Dr.Jit will attempt an implicit conversion if another type is provided. active

  • (object) – an optional 1D dynamic Dr.Jit mask array (e.g., drjit.scalar.ArrayXb or drjit.cuda.Bool) specifying active components. Dr.Jit will attempt an implicit conversion if another type is provided. The default is True.

Returns:

An instance of type dtype containing the result of the gather operation.

Return type:

object

drjit.scatter(target, value, index, active=True)#

Scatter values into a flat array or nested data structure

This operation performs a scatter (i.e., indirect memory write) of the value parameter to the target array at position index. The optional active argument can be used to disable some of the individual write operations, which is useful when not all provided values or indices are valid.

This operation can be used in the following different ways:

1. When target is a 1D Dr.Jit array like drjit.llvm.ad.Float, this operation implements a parallelized version of the Python array indexing expression target[index] = value with optional masking. Example:

target = dr.empty(dr.cuda.Float, 1024*1024)
value = dr.cuda.Float([...])
index = dr.cuda.UInt([...]) # Note: negative indices are not permitted
dr.scatter(target, value=value, index=index)

2. When target is a more complex type (e.g. a custom source structure, nested Dr.Jit array, tuple, list, dictionary, etc.), the behavior depends:

  • When target and value are of the same type, the scatter operation

    threads through entries and invokes itself recursively. For example, the scatter operation in

    target = dr.cuda.Array3f(...)
    value = dr.cuda.Array3f(...)
    index = dr.cuda.UInt([...])
    dr.scatter(target, value, index)
    

    is equivalent to

    dr.scatter(target.x, value.x, index)
    dr.scatter(target.y, value.y, index)
    dr.scatter(target.z, value.z, index)
    
  • Otherwise, the operation flattens the value array and writes it using

    C-style ordering with a suitably modified index. For example, the scatter below writes 3D vectors into a 1D array.

    target = dr.cuda.Float(...)
    value = dr.cuda.Array3f(...)
    index = dr.cuda.UInt([...])
    dr.scatter(target, value, index)
    

    and is equivalent to

    dr.scatter(target, value.x, index*3 + 0)
    dr.scatter(target, value.y, index*3 + 1)
    dr.scatter(target, value.z, index*3 + 2)
    

Danger

The indices provided to this operation are unchecked. Out-of-bounds writes are undefined behavior (if not disabled via the active parameter) and may crash the application. Negative indices are not permitted.

Dr.Jit makes no guarantees about the expected behavior when a scatter operation has conflicts, i.e., when a specific position is written multiple times by a single drjit.scatter() operation.

Parameters:
  • target (object) – The object into which data should be written (typically a 1D Dr.Jit array, but other variations are possible as well, see the description above.)

  • value (object) – The values to be written (typically of type type(target), but other variations are possible as well, see the description above.) Dr.Jit will attempt an implicit conversion if the the input is not an array type.

  • index (object) – a 1D dynamic unsigned 32-bit Dr.Jit array (e.g., drjit.scalar.ArrayXu or drjit.cuda.UInt) specifying gather indices. Dr.Jit will attempt an implicit conversion if another type is provided.

  • active (object) – an optional 1D dynamic Dr.Jit mask array (e.g., drjit.scalar.ArrayXb or drjit.cuda.Bool) specifying active components. Dr.Jit will attempt an implicit conversion if another type is provided. The default is True.

drjit.scatter_reduce(op, target, value, index, active=True)#

Perform a read-modify-write operation on a flat array or nested data structure

This function performs a read-modify-write operation of the value parameter to the target array at position index. The optional active argument can be used to disable some of the individual write operations, which is useful when not all provided values or indices are valid.

The operation to be applied is defined by tje op argument (see drjit.ReduceOp).

This operation can be used in the following different ways:

1. When target is a 1D Dr.Jit array like drjit.llvm.ad.Float, this operation implements a parallelized version of the expression target[index] = op(value, target[index]) with optional masking. Example:

target = dr.cuda.Float([...])
value = dr.cuda.Float([...])
index = dr.cuda.UInt([...]) # Note: negative indices are not permitted
dr.scatter_reduce(dr.ReduceOp.Add, target, value=value, index=index)

2. When target is a more complex type (e.g. a custom source structure, nested Dr.Jit array, tuple, list, dictionary, etc.), then target and value must be of the same type. The scatter reduce operation threads through entries and invokes itself recursively. For example, the scatter operation in

target = dr.cuda.Array3f(...)
value = dr.cuda.Array3f(...)
index = dr.cuda.UInt([...])
dr.scatter_reduce(dr.ReduceOp.Add, target, value, index)

is equivalent to

dr.scatter_reduce(dr.ReduceOp.Add, target.x, value.x, index)
dr.scatter_reduce(dr.ReduceOp.Add, target.y, value.y, index)
dr.scatter_reduce(dr.ReduceOp.Add, target.z, value.z, index)

Danger

The indices provided to this operation are unchecked. Out-of-bounds writes are undefined behavior (if not disabled via the active parameter) and may crash the application. Negative indices are not permitted.

Parameters:
  • op (drjit.ReduceOp) – Operation to be perform in the reduction.

  • target (object) – The object into which data should be written (typically a 1D Dr.Jit array, but other variations are possible as well, see the description above.)

  • value (object) – The values to be written (typically of type type(target), but other variations are possible as well, see the description above.) Dr.Jit will attempt an implicit conversion if the the input is not an array type.

  • index (object) – a 1D dynamic unsigned 32-bit Dr.Jit array (e.g., drjit.scalar.ArrayXu or drjit.cuda.UInt) specifying gather indices. Dr.Jit will attempt an implicit conversion if another type is provided.

  • active (object) – an optional 1D dynamic Dr.Jit mask array (e.g., drjit.scalar.ArrayXb or drjit.cuda.Bool) specifying active components. Dr.Jit will attempt an implicit conversion if another type is provided. The default is True.

drjit.scatter_inc(target, index, active=True)#

Atomically increment a value within an unsigned 32-bit integer array and return the value prior to the update.

This operation works just like the drjit.scatter_reduce() operation for 32-bit unsigned integer operands, but with a fixed value=1 parameter and reduce_op=ReduceOp::Add.

The main difference is that this variant additionally returns the old value of the target array prior to the atomic update in contrast to the more general scatter-reduction, which just returns None. The operation also supports masking—the return value in the unmasked case is undefined. Both target and index parameters must be 1D unsigned 32-bit arrays.

This operation is a building block for stream compaction: threads can scatter-increment a global counter to request a spot in an array and then write their result there. The recipe for this is look as follows:

ctr = UInt32(0) # Counter
active = drjit.ones(Bool, len(data_1)) # .. or a more complex condition

my_index = dr.scatter_inc(target=ctr, index=UInt32(0), mask=active)

dr.scatter(
    target=data_compact_1,
    value=data_1,
    index=my_index,
    mask=active
)

dr.scatter(
    target=data_compact_2,
    value=data_2,
    index=my_index,
    mask=active
)

When following this approach, be sure to provide the same mask value to the drjit.scatter_inc() and subsequent drjit.scatter() operations.

The function drjit.scatter_inc() exhibits the following unusual behavior compared to regular Dr.Jit operations: the return value references the instantaneous state during a potentially large sequence of atomic operations. This instantaneous state is not reproducible in later kernel evaluations, and Dr.Jit will refuse to do so when the computed index is reused. In essence, the variable is “consumed” by the process of evaluation.

my_index = dr.scatter_inc(target=ctr, index=UInt32(0), mask=active)
dr.scatter(
    target=data_compact_1,
    value=data_1,
    index=my_index,
    mask=active
)

dr.eval(data_compact_1) # Run Kernel #1

dr.scatter(
    target=data_compact_2,
    value=data_2,
    index=my_index, # <-- oops, reusing my_index in another kernel.
    mask=active     #     This raises an exception.
)

To get the above code to work, you will need to evaluate my_index at the same time to materialize it into a stored (and therefore trivially reproducible) representation. For this, ensure that the size of the active mask matches len(data_*) and that it is not the trivial True default mask (otherwise, the evaluated my_index will be scalar).

dr.eval(data_compact_1, my_index)

Such multi-stage evaluation is potentially inefficient and may defeat the purpose of performing stream compaction in the first place. In general, prefer keeping all scatter operations involving the computed index in the same kernel, and then this issue does not arise.

The implementation of drjit.scatter_inc() performs a local reduction first, followed by a single atomic write per SIMD packet/warp. This is done to reduce contention from a potentially very large number of atomic operations targeting the same memory address. Fully masked updates do not cause memory traffic.

There is some conceptual overlap between this function and drjit.compress(), which can likewise be used to reduce a stream to a smaller subset of active items. The downside of drjit.compress() is that it requires evaluating the variables to be reduced, which can be very costly in terms of of memory traffic and storage footprint. Reducing through drjit.scatter_inc() does not have this limitation: it can operate on symbolic arrays that greatly exceed the available device memory. One advantage of drjit.compress() is that it essentially boils down to a realtively simple prefix sum, which does not require atomic memory operations (these can be slow in some cases).

drjit.ravel(array, order='A')#

Convert the input into a contiguous flat array

This operation takes a Dr.Jit array, typically with some static and some dynamic dimensions (e.g., drjit.cuda.Array3f with shape 3xN), and converts it into a flattened 1D dynamically sized array (e.g., drjit.cuda.Float) using either a C or Fortran-style ordering convention.

It can also convert Dr.Jit tensors into a flat representation, though only C-style ordering is supported in this case.

For example,

x = dr.cuda.Array3f([1, 2], [3, 4], [5, 6])
y = dr.ravel(x, order=...)

will produce

  • [1, 3, 5, 2, 4, 6] with order='F' (the default for Dr.Jit arrays),

which means that X/Y/Z components alternate. - [1, 2, 3, 4, 5, 6] with order='C', in which case all X coordinates are written as a contiguous block followed by the Y- and then Z-coordinates.

Danger

Currently C-style ordering is not implemented for tensor types.

Parameters:
  • array (drjit.ArrayBase) – An arbitrary Dr.Jit array or tensor

  • order (str) – A single character indicating the index order. 'F' indicates column-major/Fortran-style ordering, in which case the first index changes at the highest frequency. The alternative 'C' specifies row-major/C-style ordering, in which case the last index changes at the highest frequency. The default value 'A' (automatic) will use F-style ordering for arrays and C-style ordering for tensors.

Returns:

A dynamic 1D array containing the flattened representation of

array with the desired ordering. The type of the return value depends on the type of the input. When array is already contiguous/flattened, this function returns it without making a copy.

Return type:

object

drjit.unravel(dtype, array, order='F')#

Load a sequence of Dr.Jit vectors/matrices/etc. from a contiguous flat array

This operation implements the inverse of drjit.ravel(). In contrast to drjit.ravel(), it requires one additional parameter (dtype) specifying type of the return value. For example,

x = dr.cuda.Float([1, 2, 3, 4, 5, 6])
y = dr.unravel(dr.cuda.Array3f, x, order=...)

will produce an array of two 3D vectors with different contents depending on the indexing convention:

  • [1, 2, 3] and [4, 5, 6] when unraveled with order='F' (the default for Dr.Jit arrays), and

  • [1, 3, 5] and [2, 4, 6] when unraveled with order='C'

Parameters:
  • dtype (type) – An arbitrary Dr.Jit array type

  • array (drjit.ArrayBase) – A dynamically sized 1D Dr.Jit array instance that is compatible with dtype. In other words, both must have the same underlying scalar type and be located imported in the same package (e.g., drjit.llvm.ad).

  • order (str) – A single character indicating the index order. 'F' (the default) indicates column-major/Fortran-style ordering, in which case the first index changes at the highest frequency. The alternative 'C' specifies row-major/C-style ordering, in which case the last index changes at the highest frequency.

Returns:

An instance of type dtype containing the result of the unravel

operation.

Return type:

object

drjit.slice(value, index=None, return_type=None)#

Slice a structure of arrays to return a single entry for a given index

This function is the equivalent to __getitem__(index) for the dynamic dimension of a Dr.Jit array or custom source structure. It can be used to access a single element out a structure of arrays for a given index.

The returned object type will differ from the type of the input value as its dynamic dimension will be removed. For static arrays (e.g. drjit.cuda.Array3f) the function will return a Python list. For custom source structure the returned type needs to be specified through the argument return_type.

Parameters:
  • value (object) – A dynamically sized 1D Dr.Jit array instance that is compatible with dtype. In other words, both must have the same underlying scalar type and be located imported in the same package (e.g., drjit.llvm.ad).

  • index (int) – Index of the entry to be returned in the structure of arrays. When not specified (or None), the provided object must have a dynamic width of 1 and this function will remove the dynamic dimension to this object by casting it into the appropriate type.

  • return_type (type) – A return type must be specified when slicing through a custom source structure. Otherwise set to None.

Returns:

Single entry of the structure of arrays.

Return type:

object

drjit.min(arg, /) float | int | drjit.ArrayBase#

Compute the minimum value in the provided input.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Minimum of the input

drjit.max(arg, /) float | int | drjit.ArrayBase#

Compute the maximum value in the provided input.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Maximum of the input

drjit.sum(arg, /) float | int | drjit.ArrayBase#

Compute the sum of all array elements.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Sum of the input

drjit.prod(arg, /) float | int | drjit.ArrayBase#

Compute the product of all array elements.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Product of the input

drjit.dot(arg0, arg1, /) float | int | drjit.ArrayBase#

Computes the dot product of two arrays.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:
Returns:

Dot product of inputs

drjit.norm(arg, /) float | int | drjit.ArrayBase#

Computes the norm of an array.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (list | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Norm of the input

drjit.shuffle(perm, value)#

Permute the entries of the provided Dr.Jit static array for the indices given in perm.

The pseudocode for this operation is

out = [value[p] for p in perm]
Parameters:
Returns:

Shuffled input array

drjit.compress(arg, /) int | drjit.ArrayBase#

Compress a mask into an array of nonzero indices.

This function takes an boolean array as input and then returns an unsigned 32-bit integer array containing the indices of nonzero entries.

It can be used to reduce a stream to a subset of active entries via the following recipe:

There is some conceptual overlap between this function and drjit.cscatter_inc(), which can likewise be used to reduce a stream to a smaller subset of active items. Please see the documentation of t drjit.cscatter_inc() for details.

Danger

This function internally performs a synchronization step.

Parameters:

arg (bool | drjit.ArrayBase) – A Python or Dr.Jit boolean type

Returns:

Array of nonzero indices

drjit.count(arg, /) int | drjit.ArrayBase#

Efficiently computes the number of entries whose boolean values are True, i.e.

(value[0] ? 1 : 0) + ... (value[Size - 1] ? 1 : 0)

For 1D arrays, count() returns a result of type int. For multidimensional arrays, the horizontal reduction is performed over the outermost dimension.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit boolean type

Returns:

Number of entries whose mask bits are turned on

drjit.block_sum(value, size)#

Sum over elements within blocks

This function adds all elements of contiguous blocks of size size in the input array value and writes them to the returned array. For example, a, b, c, d, e, f turns into a+b, c+d, e+f when size == 2. The length of the input array must be a multiple of size.

Parameters:
  • arg (drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

  • size (int) – size of the block

Returns:

Sum over elements within blocks

Mask operations#

drjit.select(condition, x, y, /)#

Select elements from inputs based on a condition

This function implements the component-wise operation

\[\]
mathrm{result}_i = begin{cases}

x_i,quad&text{if condition}_i,\ y_i,quad&text{otherwise.}

end{cases}

Parameters:
Returns:

Component-wise result of the selection operation

Return type:

float | int | drjit.ArrayBase

drjit.all(arg, /) bool | drjit.ArrayBase#

Computes whether all input elements evaluate to True.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Boolean array

drjit.any(arg, /) bool | drjit.ArrayBase#

Computes whether any of the input elements evaluate to True.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Boolean array

drjit.all_nested(arg, /) bool#

Iterates all() until the type of the return value no longer changes. This can be used to reduce a nested mask array into a single value.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Boolean

drjit.any_nested(arg, /) bool#

Iterates any() until the type of the return value no longer changes. This can be used to reduce a nested mask array into a single value.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Boolean

drjit.eq(a, b)#

Return the element-wise comparison of the two arguments

This function falls back to == when none of the arguments are Dr.Jit arrays.

Parameters:
  • a (object) – Input array.

  • b (object) – Input array.

Returns:

Output array, element-wise comparison of a and b

Return type:

object

drjit.neq(a, b)#

Return the element-wise not-equal comparison of the two arguments

This function falls back to != when none of the arguments are Dr.Jit arrays.

Parameters:
  • a (object) – Input array.

  • b (object) – Input array.

Returns:

Output array, element-wise comparison of a != b

Return type:

object

drjit.isinf(arg, /)#

Performs an element-wise test for positive or negative infinity

Parameters:

arg (object) – A Dr.Jit array or other kind of numeric sequence type.

Returns:

A mask value describing the result of the test

Return type:

mask_t(arg)

drjit.isnan(arg, /)#

Performs an element-wise test for NaN (Not a Number) values

Parameters:

arg (object) – A Dr.Jit array or other kind of numeric sequence type.

Returns:

A mask value describing the result of the test.

Return type:

mask_t(arg)

drjit.isfinite(arg, /)#

Performs an element-wise test that checks whether values are finite and not equal to NaN (Not a Number)

Parameters:

arg (object) – A Dr.Jit array or other kind of numeric sequence type.

Returns:

A mask value describing the result of the test

Return type:

mask_t(arg)

drjit.allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)#

Returns True if two arrays are element-wise equal within a given error tolerance.

The function considers both absolute and relative error thresholds. Specifically a and b are considered equal if all elements satisfy

\[|a - b| \le |b| \cdot \mathrm{rtol} + \mathrm{atol}. \]
Parameters:
  • a (object) – A Dr.Jit array or other kind of numeric sequence type.

  • b (object) – A Dr.Jit array or other kind of numeric sequence type.

  • rtol (float) – A relative error threshold. The default is \(10^{-5}\).

  • atol (float) – An absolute error threshold. The default is \(10^{-8}\).

  • equal_nan (bool) – If a and b both contain a NaN (Not a Number) entry, should they be considered equal? The default is False.

Returns:

The result of the comparison.

Return type:

bool

Miscellaneous operations#

drjit.shape(arg, /)#

Return a tuple describing dimension and shape of the provided Dr.Jit array or tensor.

When the arrays is ragged, the implementation signals a failure by returning None. A ragged array has entries of incompatible size, e.g. [[1, 2], [3, 4, 5]]. Note that an scalar entries (e.g. [[1, 2], [3]]) are acceptable, since broadcasting can effectively convert them to any size.

The expressions drjit.shape(arg) and arg.shape are equivalent.

Parameters:

arg (drjit.ArrayBase) – An arbitrary Dr.Jit array or tensor

Returns:

A tuple describing the dimension and shape of the

provided Dr.Jit input array or tensor. When the input array is ragged (i.e., when it contains components with mismatched sizes), the function returns None.

Return type:

tuple | NoneType

drjit.width(arg, /)#

Return the width of the provided dynamic Dr.Jit array, tensor, or custom data structure.

The function returns 1 if the input variable isn’t a Dr.Jit array, tensor, or custom data structure.

Parameters:

arg (drjit.ArrayBase) – An arbitrary Dr.Jit array, tensor, or custom data structure

Returns:

The dynamic width of the provided variable.

Return type:

int

drjit.resize(arg, size)#

Resize in-place the provided Dr.Jit array, tensor, or custom data structure to a new size.

The provided variable must have a size of zero or one originally otherwise this function will fail.

When the provided variable doesn’t have a size of 1 and its size exactly matches size the function does nothing. Otherwise, it fails.

Parameters:
  • arg (drjit.ArrayBase) – An arbitrary Dr.Jit array, tensor, or custom data structure to be resized

  • size (int) – The new size

Perform binary search over a range given a predicate pred, which monotonically decreases over this range (i.e. max one True -> False transition).

Given a (scalar) start and end index of a range, this function evaluates a predicate floor(log2(end-start) + 1) times with index values on the interval [start, end] (inclusive) to find the first index that no longer satisfies it. Note that the template parameter Index is automatically inferred from the supplied predicate. Specifically, the predicate takes an index array as input argument. When pred is False for all entries, the function returns start, and when it is True for all cases, it returns end.

The following code example shows a typical use case: data contains a sorted list of floating point numbers, and the goal is to map floating point entries of x to the first index j such that data[j] >= threshold (and all of this of course in parallel for each vector element).

dtype = dr.llvm.Float
data = dtype(...)
threshold = dtype(...)

index = dr.binary_search(
    0, len(data) - 1,
    lambda index: dr.gather(dtype, data, index) < threshold
)
Parameters:
  • start (int) – Starting index for the search range

  • end (int) – Ending index for the search range

  • pred (function) – The predicate function to be evaluated

Returns:

Index array resulting from the binary search

drjit.upsample(source, shape=None, scale_factor=None, align_corners=False)#

Up-sample the input tensor or texture according to the provided shape.

Alternatively to specifying the target shape, a scale factor can be provided.

The behavior of this function depends on the type of source:

1. When source is a Dr.Jit tensor, nearest neighbor up-sampling will use hence the target shape values must be multiples of the source shape values. When scale_factor is used, its values must be integers.

2. When source is a Dr.Jit texture type, the up-sampling will be performed according to the filter mode set on the input texture. Target shape values are not required to be multiples of the source shape values. When scale_factor is used, its values must be integers.

Parameters:
  • source (object) – A Dr.Jit tensor or texture type.

  • shape (list) – The target shape (optional)

  • scale_factor (list) – The scale factor to apply to the current shape

  • (optional)

  • align_corners (bool) – Defines whether or not the corner pixels of the

  • the (input and output should be aligned. This allows the values at) –

  • is (a Dr.Jit texture type performing linear interpolation. The default) –

  • is

  • False.

Returns:

the up-sampled tensor or texture object. The type of the output will be the same as the type of the source.

Return type:

object

drjit.tzcnt(arg, /)#

Return the number of trailing zero bits.

This function assumes that arg is an integer array.

Parameters:

arg (int | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

number of trailing zero bits in the input array

Return type:

int | drjit.ArrayBase

drjit.lzcnt(arg, /)#

Return the number of leading zero bits.

This function assumes that arg is an integer array.

Parameters:

arg (int | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

number of leading zero bits in the input array

Return type:

int | drjit.ArrayBase

drjit.popcnt(arg, /)#

Return the number of nonzero zero bits.

This function assumes that arg is an integer array.

Parameters:

arg (int | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

number of nonzero zero bits in the input array

Return type:

int | drjit.ArrayBase

Function dispatching#

drjit.switch(indices, funcs, *args)#

Dispatches a call to one of the given functions based on the given indices.

def f1(x):
     return x * 10

def f2(x):
     return x * 100

arg = dr.llvm.Float([1.0, 2.0, 3.0, 4.0])
indices = dr.llvm.UInt32([0, 1, 1, 0])

res = dr.switch(indices, [f1, f2], arg)

# [10.0, 200.0, 300.0, 40.0]
Parameters:
  • indices (drjit.ArrayBase) – a list of indices to choose the functions

  • funcs (list) – a list of functions to dispatch based on indices

  • args (tuple) – the arguments to pass to the functions

Returns:

the result of the function call dispatched based on the indices

Return type:

object

drjit.dispatch(instances, func, *args)#

Dispatches a call to the given function on an array of instances of a class registered to the Dr.Jit registry. The provided function is assumed to take at least one argument that will represent the individual instance of the class this function is recorded on.

This routine can be used to perform dynamic dispatch to Python methods of instances. The code generation underlying Dr.Jit will transform arbitrary sequences of such method calls into unified functions in the generated code. In other words, a single dispatch call that calls two methods is potentially significantly faster than two separate dispatches.

pointers = ... # Dr.Jit pointer array type

def func(self, arg):
    v = self.foo() # call foo() on a single instance at a time
    return v + arg

arg = dr.llvm.Float([1.0, 2.0, 3.0, 4.0])

res = dr.dispatch(pointers, func, arg)
Parameters:
  • instances (drjit.ArrayBase) – array of pointers to instances to dispatch the given function on.

  • func (dict) – function to dispatch on all instances.

  • args (tuple) – the arguments to pass to the function.

Returns:

the result of the function evaluated for the given instances.

Return type:

object

Just-in-time compilation#

drjit.schedule(*args)#

Schedule the provided JIT variable(s) for later evaluation

This function causes args to be evaluated by the next kernel launch. In other words, the effect of this operation is deferred: the next time that Dr.Jit’s LLVM or CUDA backends compile and execute code, they will include the trace of the specified variables in the generated kernel and turn them into an explicit memory-based representation.

Scheduling and evaluation of traced computation happens automatically, hence it is rare that a user would need to call this function explicitly. Explicit scheduling can improve performance in certain cases—for example, consider the following code:

# Computation that produces Dr.Jit arrays
a, b = ...

# The following line launches a kernel that computes 'a'
print(a)

# The following line launches a kernel that computes 'b'
print(b)

If the traces of a and b overlap (perhaps they reference computation from an earlier step not shown here), then this is inefficient as these steps will be executed twice. It is preferable to launch bigger kernels that leverage common subexpressions, which is what drjit.schedule() enables:

a, b = ... # Computation that produces Dr.Jit arrays

# Schedule both arrays for deferred evaluation, but don't evaluate yet
dr.schedule(a, b)

# The following line launches a kernel that computes both 'a' and 'b'
print(a)

# References the stored array, no kernel launch
print(b)

Note that drjit.eval() would also have been a suitable alternative in the above example; the main difference to drjit.schedule() is that it does the evaluation immediately without deferring the kernel launch.

This function accepts a variable-length keyword argument and processes it as follows:

  • It recurses into sequences (tuple, list, etc.)

  • It recurses into the values of mappings (dict, etc.)

  • It recurses into the fields of custom data structures.

During recursion, the function gathers all unevaluated Dr.Jit arrays. Evaluated arrays and incompatible types are ignored. Multiple variables can be equivalently scheduled with a single drjit.schedule() call or a sequence of calls to drjit.schedule(). Variables that are garbage collected between the original drjit.schedule() call and the next kernel launch are ignored and will not be stored in memory.

Parameters:

*args (tuple) – A variable-length list of Dr.Jit array instances, custom data structures, sequences, or mappings. The function will recursively traverse data structures to discover all Dr.Jit arrays.

Returns:

True if a variable was scheduled, False if the operation did not do anything.

Return type:

bool

drjit.eval(*args)#

Immediately evaluate the provided JIT variable(s)

This function immediately invokes Dr.Jit’s LLVM or CUDA backends to compile and then execute a kernel containing the trace of the specified variables, turning them into an explicit memory-based representation. The generated kernel(s) will also include previously scheduled computation. The function drjit.eval() internally calls drjit.schedule()—specifically,

dr.eval(arg_1, arg_2, ...)

is equivalent to

dr.schedule(arg_1, arg_2, ...)
dr.eval()

Variable evaluation happens automatically as needed, hence it is rare that a user would need to call this function explicitly. Explicit evaluation can slightly improve performance in certain cases (the documentation of drjit.schedule() shows an example of such a use case.)

This function accepts a variable-length keyword argument and processes it as follows:

  • It recurses into sequences (tuple, list, etc.)

  • It recurses into the values of mappings (dict, etc.)

  • It recurses into the fields of custom data structures.

During recursion, the function gathers all unevaluated Dr.Jit arrays. Evaluated arrays and incompatible types are ignored.

Parameters:

*args (tuple) – A variable-length list of Dr.Jit array instances, custom data structures, sequences, or mappings. The function will recursively traverse data structures to discover all Dr.Jit arrays.

Returns:

True if a variable was evaluated, False if the operation did not do anything.

Return type:

bool

drjit.printf_async(fmt, *args, active=True)#

Print the specified variable contents from the kernel asynchronously.

This function inserts a print statement directly into the kernel being generated. Note that this may produce a very large volume of output, and a nonzero active parameter can be supplied to suppress it based on condition.

Parameters:
  • fmt (str) – The string to be printed. It might contain format specifiers (e.g. subsequences beginning with %)

  • *args (tuple) – Additional array arguments to be formatted and inserted in the printed string replacing their respective specifiers.

  • active (bool | drjit.ArrayBase) – Mask array to suppress printing specific elements in the supplied additional arrays.

drjit.graphviz(as_str=False)#

Assembles a graphviz diagram for the computational graph trace by the JIT.

Parameters:

as_str (bool) – whether the function should return the graphviz object as a string representation or not.

Returns:

the graphviz obj (or its string representation).

Return type:

object

drjit.label(arg)#

Returns the label of a given Dr.Jit array.

Parameters:

arg (object) – a Dr.Jit array instance.

Returns:

the label of the given variable.

Return type:

str

drjit.set_label(*args, **kwargs)#

Sets the label of a provided Dr.Jit array, either in the JIT or the AD system.

When a custom data structure is provided, the field names will be used as suffix for the variables labels.

When a sequence or static array is provided, the item’s indices will be appended to the label.

When a mapping is provided, the item’s key will be appended to the label.

Parameters:
  • *arg (tuple) – a Dr.Jit array instance and its corresponding label str value.

  • **kwarg (dict) – A set of (keyword, object) pairs.

Type traits#

The functions in this section can be used to infer properties or types of Dr.Jit arrays.

The naming convention with a trailing _v or _t indicates whether a function returns a value or a type. Evaluation takes place at runtime within Python. In C++, these expressions are all constexpr (i.e., evaluated at compile time.).

Array type tests#

drjit.is_array_v(arg, /)#

Check if the input is a Dr.Jit array instance or type

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg or type(arg) is a Dr.Jit array type, and False otherwise

Return type:

bool

drjit.is_mask_v(arg, /)#

Check whether the input array instance or type is a Dr.Jit mask array or a Python bool value/type.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents a Dr.Jit mask array or Python bool instance or type.

Return type:

bool

drjit.is_float_v(arg, /)#

Check whether the input array instance or type is a Dr.Jit floating point array or a Python float value/type.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents a Dr.Jit floating point array or Python float instance or type.

Return type:

bool

drjit.is_integral_v(arg, /)#

Check whether the input array instance or type is an integral Dr.Jit array or a Python int value/type.

Note that a mask array is not considered to be integral.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an integral Dr.Jit array or Python int instance or type.

Return type:

bool

drjit.is_arithmetic_v(arg, /)#

Check whether the input array instance or type is an arithmetic Dr.Jit array or a Python int or float value/type.

Note that a mask type (e.g. bool, drjit.scalar.Array2b, etc.) is not considered to be arithmetic.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an arithmetic Dr.Jit array or Python int or float instance or type.

Return type:

bool

drjit.is_signed_v(arg, /)#

Check whether the input array instance or type is an signed Dr.Jit array or a Python int or float value/type.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an signed Dr.Jit array or Python int or float instance or type.

Return type:

bool

drjit.is_unsigned_v(arg, /)#

Check whether the input array instance or type is an unsigned integer Dr.Jit array or a Python bool value/type (masks and boolean values are also considered to be unsigned).

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an unsigned Dr.Jit array or Python bool instance or type.

Return type:

bool

drjit.is_jit_v(arg, /)#

Check whether the input array instance or type represents a type that undergoes just-in-time compilation.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an array type from the drjit.cuda.* or drjit.llvm.* namespaces, and False otherwise.

Return type:

bool

drjit.is_llvm_v(arg, /)#

Check whether the input is a Dr.Jit LLVM array instance or type.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an array type from the drjit.llvm.* namespace, and False otherwise.

Return type:

bool

drjit.is_diff_v(arg, /)#

Check whether the input is a differentiable Dr.Jit array instance or type.

Note that this is a type-based statement that is unrelated to mathematical differentiability. For example, the integral type drjit.cuda.ad.Int from the CUDA AD namespace satisfies is_diff_v(..) = 1.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg represents an array type from the drjit.[cuda/llvm].ad.* namespace, and False otherwise.

Return type:

bool

drjit.is_complex_v(arg, /)#

Check whether the input is a Dr.Jit array instance or type representing a complex number.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if the test was successful, and False otherwise.

Return type:

bool

drjit.is_matrix_v(arg, /)#

Check whether the input is a Dr.Jit array instance or type representing a matrix.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if the test was successful, and False otherwise.

Return type:

bool

drjit.is_quaternion_v(arg, /)#

Check whether the input is a Dr.Jit array instance or type representing a quaternion.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if the test was successful, and False otherwise.

Return type:

bool

drjit.is_tensor_v(arg, /)#

Check whether the input is a Dr.Jit array instance or type representing a tensor.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if the test was successful, and False otherwise.

Return type:

bool

drjit.is_special_v(arg, /)#

Check whether the input is a special Dr.Jit array instance or type.

A special array type requires precautions when performing arithmetic operations like multiplications (complex numbers, quaternions, matrices).

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if the test was successful, and False otherwise.

Return type:

bool

drjit.is_struct_v(arg, /)#

Check if the input is a Dr.Jit-compatible data structure

Custom data structures can be made compatible with various Dr.Jit operations by specifying a DRJIT_STRUCT member. See the section on custom data structure for details. This type trait can be used to check for the existence of such a field.

Parameters:

arg (object) – An arbitrary Python object

Returns:

True if arg has a DRJIT_STRUCT member

Return type:

bool

Array shape#

drjit.size_v(arg, /)#

Return the (static) size of the outermost dimension of the provided Dr.Jit array instance or type

Note that this function mainly exists to query type-level information. Use the Python len() function to query the size in a way that does not distinguish between static and dynamic arrays.

Parameters:

arg (object) – An arbitrary Python object

Returns:

Returns either the static size or drjit.Dynamic when arg is a dynamic Dr.Jit array. Returns 1 for all other types.

Return type:

int

drjit.depth_v(arg, /)#

Return the depth of the provided Dr.Jit array instance or type

For example, an array consisting of floating point values (for example, drjit.scalar.Array3f) has depth 1, while an array consisting of sub-arrays (e.g., drjit.cuda.Array3f) has depth 2.

Parameters:

arg (object) – An arbitrary Python object

Returns:

Returns the depth of the input, if it is a Dr.Jit array instance or type. Returns 0 for all other inputs.

Return type:

int

drjit.Dynamic: int = -1#

Special size value used to identify dynamic arrays in array_size_v().

Others#

drjit.reinterpret_array_v(dtype, value)#

Reinterpret the provided Dr.Jit array/tensor into a specified type.

Parameters:
  • dtype (type) – Target type for the reinterpretation.

  • value (object) – Dr.Jit array or tensor to reinterpret.

Returns:

Result of the conversion as described above.

Return type:

object

Standard mathematical functions#

drjit.abs(arg, /)#

Compute the absolute value of the provided input.

Parameters:

arg (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Absolute value of the input)

Return type:

float | int | drjit.ArrayBase

drjit.minimum(arg0, arg1, /) float | int | drjit.ArrayBase#

Compute the element-wise minimum value of the provided inputs.

This function returns a result of the type type(arg0 + arg1) (i.e., according to the usual implicit type conversion rules).

Parameters:
  • arg0 (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

  • arg1 (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Minimum of the input(s)

drjit.maximum(arg0, arg1, /) float | int | drjit.ArrayBase#

Compute the element-wise maximum value of the provided inputs.

This function returns a result of the type type(arg0 + arg1) (i.e., according to the usual implicit type conversion rules).

Parameters:
  • arg0 (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

  • arg1 (float | int | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Maximum of the input(s)

drjit.clip(value, min, max, /)#

Clip the provided input to the given interval.

This function is equivalent to

dr.maximum(dr.minimum(value, max), min)

Dr.Jit also defines drjit.clamp() as an alias of this function.

Parameters:
Returns:

Clipped input

Return type:

float | drjit.ArrayBase

drjit.clamp(value, min, max, /)#

Clip the provided input to the given interval.

This function is equivalent drjit.clamp().

Parameters:
Returns:

Clipped input

Return type:

float | drjit.ArrayBase

drjit.fma(arg0, arg1, arg2, /)#

Perform a fused multiply-add (FMA) operation.

Given arguments arg0, arg1, and arg2, this operation computes arg0 * arg1 + arg2 using only one final rounding step. The operation is not only more accurate, but also more efficient, since FMA maps to a native machine instruction on platforms targeted by Dr.Jit.

While FMA is traditionally a floating point operation, Dr.Jit also implements FMA for integer arrays and maps it onto dedicated instructions provided by the backend if possible (e.g. mad.lo.* for CUDA/PTX).

Parameters:
Returns:

Result of the FMA operation

Return type:

float | drjit.ArrayBase

drjit.ceil(arg, /)#

Evaluate the ceiling, i.e. the smallest integer >= arg.

The function does not convert the type of the input array. A separate cast is necessary when integer output is desired.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Ceiling of the input

Return type:

float | drjit.ArrayBase

drjit.floor(arg, /)#

Evaluate the floor, i.e. the largest integer <= arg.

The function does not convert the type of the input array. A separate cast is necessary when integer output is desired.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Floor of the input

Return type:

float | drjit.ArrayBase

drjit.trunc(arg, /)#

Truncates arg to the nearest integer by towards zero.

The function does not convert the type of the input array. A separate cast is necessary when integer output is desired.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Truncated result

Return type:

float | drjit.ArrayBase

drjit.round(arg, /)#

Rounds arg to the nearest integer using Banker’s rounding for half-way values.

This function is equivalent to std::rint in C++. It does not convert the type of the input array. A separate cast is necessary when integer output is desired.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Rounded result

Return type:

float | drjit.ArrayBase

drjit.sqrt(arg, /)#

Evaluate the square root of the provided input.

Negative inputs produce a NaN output value.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Square root of the input

Return type:

float | drjit.ArrayBase

drjit.cbrt(arg, /)#

Evaluate the cube root of the provided input.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Cube root of the input

Return type:

float | drjit.ArrayBase

drjit.rcp(arg, /)#

Evaluate the reciprocal (1 / arg) of the provided input.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”). The result is slightly approximate in this case (refer to the documentation of the instruction rcp.approx.ftz.f32 in the NVIDIA PTX manual for details).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Reciprocal of the input

Return type:

float | drjit.ArrayBase

drjit.rsqrt(arg, /)#

Evaluate the reciprocal square root (1 / sqrt(arg)) of the provided input.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”). The result is slightly approximate in this case (refer to the documentation of the instruction rsqrt.approx.ftz.f32 in the NVIDIA PTX manual for details).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Reciprocal square root of the input

Return type:

float | drjit.ArrayBase

drjit.lerp(a, b, t, /)#

Blends between the values a and b using the expression \(a \cdot (1 - t) + b \cdot t\)

Parameters:
Returns:

Blended value

Return type:

float | int | drjit.ArrayBase

drjit.normalize(arg, /) drjit.ArrayBase#

Normalizes the provided array.

When the argument is a dynamic array, function performs a horizontal reduction. Please see the section on horizontal reductions for details.

Parameters:

arg (list | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Normalized input array

drjit.log2i(arg, /)#

Return the floor of the base-two logarithm.

This function assumes that arg is an integer array.

Parameters:

arg (int | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

number of leading zero bits in the input array

Return type:

int | drjit.ArrayBase

drjit.erf(arg, /)#

Evaluates the error function defined as

\[\mathrm{erf}(x)=\frac{2}{\sqrt{\pi}}\int_0^x e^{-t^2}\,\mathrm{d}t.\]

Requires a real-valued input array x.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Value of the error function at the input value

Return type:

float | drjit.ArrayBase

drjit.erfinv(arg, /)#

Evaluates the inverse of the error function drjit.erf().

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Value of the inverse of the error function at the input value

Return type:

float | drjit.ArrayBase

drjit.lgamma(arg, /)#

Evaluates the natural logarithm of the Gamma function.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Value of the natural logarithm of the Gamma function at the input value

Return type:

float | drjit.ArrayBase

drjit.tgamma(arg, /)#

Evaluates the Gamma function defined as

\[\Gamma(x)=\int_0^\infty t^{x-1} e^{-t}\,\mathrm{d}t.\]
Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Value of the Gamma function at the input value

Return type:

float | drjit.ArrayBase

drjit.hypot(a, b)#

Computes \(\sqrt{x^2+y^2}\) while avoiding overflow and underflow.

Parameters:

arg (list | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type

Returns:

Hypotenuse value

drjit.sign(arg, /)#

Return the element-wise sign of the provided array.

Parameters:

arg (int | float | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

Sign of the input array

Return type:

float | int | drjit.ArrayBase

drjit.copysign(arg0, arg1, /)#

Copy the sign of arg1 to ``arg0` element-wise.

Parameters:
  • arg0 (int | float | drjit.ArrayBase) – A Python or Dr.Jit array to change the sign of

  • arg1 (int | float | drjit.ArrayBase) – A Python or Dr.Jit array to copy the sign from

Returns:

The values of arg0 with the sign of arg1

Return type:

float | int | drjit.ArrayBase

drjit.mulsign(arg0, arg1, /)#

Multiply arg0 by the sign of ``arg1` element-wise.

This function is equivalent to

a * dr.sign(b)
Parameters:
  • arg0 (int | float | drjit.ArrayBase) – A Python or Dr.Jit array to multiply the sign of

  • arg1 (int | float | drjit.ArrayBase) – A Python or Dr.Jit array to take the sign from

Returns:

The values of arg0 multiplied with the sign of arg1

Return type:

float | int | drjit.ArrayBase

drjit.arg(arg, /)#

Return the argument of a complex Dr.Jit array.

When the provided array isn’t an instance of drjit.Complex, this function assumes that the input array represents the real part of a complex variable.

Parameters:

arg (int | float | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

Argument of the complex input array

Return type:

float | drjit.ArrayBase

drjit.real(arg, /)#

Return the real part of a complex Dr.Jit array.

When the provided array isn’t an instance of drjit.Complex or drjit.Quaternion, this function returns the input unchanged.

Parameters:

arg (int | float | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

Real part of the input array

Return type:

float | drjit.ArrayBase

drjit.imag(arg, /)#

Return the imaginary part of a complex Dr.Jit array.

When the provided array isn’t an instance of drjit.Complex or drjit.Quaternion, this function returns the input unchanged.

Parameters:

arg (int | float | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

Imaginary part of the input array

Return type:

float | drjit.ArrayBase

drjit.conj(arg, /)#

Return the complex conjugate of a provided Dr.Jit array.

When the provided array isn’t an instance of drjit.Complex or drjit.Quaternion, this function returns the input unchanged.

Parameters:

arg (int | float | drjit.ArrayBase) – A Python or Dr.Jit array

Returns:

Real part of the input array

Return type:

float | drjit.ArrayBase

drjit.cross(a, b, /)#

Returns the cross-product of the two input 3D arrays

Parameters:
Returns:

Cross-product of the two input 3D arrays

drjit.sh_eval(arg, order, /)#

Evaluates the real spherical harmonics basis functions up to and including order order.

The directions provided to sh_eval must be normalized 3D vectors (i.e. using Cartesian instead of spherical coordinates).

This function supports evaluation order up to 10 (e.g. order=9).

Parameters:
  • arg (drjit.ArrayBase) – A 3D Dr.Jit array type for the direction to be evaluated

  • order (int) – Order of the spherical harmonic evaluation

Returns:

List of spherical harmonics coefficients

Return type:

list

Transcendental functions#

Dr.Jit implements the most common transcendental functions using methods that are based on the CEPHES math library. The accuracy of these approximations is documented in a set of tables below.

Trigonometric functions#

drjit.sin(arg, /)#

Sine approximation based on the CEPHES library.

The implementation of this function is designed to achieve low error on the domain \(|x| < 8192\) and will not perform as well beyond this range. See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Sine of the input

Return type:

float | drjit.ArrayBase

drjit.cos(arg, /)#

Cosine approximation based on the CEPHES library.

The implementation of this function is designed to achieve low error on the domain \(|x| < 8192\) and will not perform as well beyond this range. See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Cosine of the input

Return type:

float | drjit.ArrayBase

drjit.sincos(arg, /)#

Sine/cosine approximation based on the CEPHES library.

The implementation of this function is designed to achieve low error on the domain \(|x| < 8192\) and will not perform as well beyond this range. See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using two operations involving the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Sine and cosine of the input

Return type:

(float, float) | (drjit.ArrayBase, drjit.ArrayBase)

drjit.tan(arg, /)#

Tangent approximation based on the CEPHES library.

The implementation of this function is designed to achieve low error on the domain \(|x| < 8192\) and will not perform as well beyond this range. See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Tangent of the input

Return type:

float | drjit.ArrayBase

drjit.asin(arg, /)#

Arcsine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Arcsine of the input

Return type:

float | drjit.ArrayBase

drjit.acos(arg, /)#

Arccosine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Arccosine of the input

Return type:

float | drjit.ArrayBase

drjit.atan(arg, /)#

Arctangent approximation

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Arctangent of the input

Return type:

float | drjit.ArrayBase

drjit.atan2(y, x, /)#

Arctangent of two values

See the section on transcendental function approximations for details regarding accuracy.

Parameters:
Returns:

Arctangent of y/x, using the argument signs to determine the quadrant of the return value

Return type:

float | drjit.ArrayBase

Hyperbolic functions#

drjit.sinh(arg, /)#

Hyperbolic sine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic sine of the input

Return type:

float | drjit.ArrayBase

drjit.cosh(arg, /)#

Hyperbolic cosine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic cosine of the input

Return type:

float | drjit.ArrayBase

drjit.sincosh(arg, /)#

Hyperbolic sine/cosine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic sine and cosine of the input

Return type:

(float, float) | (drjit.ArrayBase, drjit.ArrayBase)

drjit.tanh(arg, /)#

Hyperbolic tangent approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic tangent of the input

Return type:

float | drjit.ArrayBase

drjit.asinh(arg, /)#

Hyperbolic arcsine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic arcsine of the input

Return type:

float | drjit.ArrayBase

drjit.acosh(arg, /)#

Hyperbolic arccosine approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic arccosine of the input

Return type:

float | drjit.ArrayBase

drjit.atanh(arg, /)#

Hyperbolic arctangent approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Hyperbolic arctangent of the input

Return type:

float | drjit.ArrayBase

Exponentials, logarithms, power function#

drjit.log2(arg, /)#

Base-2 exponential approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Base-2 logarithm of the input

Return type:

float | drjit.ArrayBase

drjit.log(arg, /)#

Natural exponential approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Natural logarithm of the input

Return type:

float | drjit.ArrayBase

drjit.exp2(arg, /)#

Base-2 exponential approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Base-2 exponential of the input

Return type:

float | drjit.ArrayBase

drjit.exp(arg, /)#

Natural exponential approximation based on the CEPHES library.

See the section on transcendental function approximations for details regarding accuracy.

When arg is a CUDA single precision array, the operation is implemented using the native multi-function unit (“MUFU”).

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Natural exponential of the input

Return type:

float | drjit.ArrayBase

drjit.power(a, b)#

Raise the first input value to the power given as second input value.

This function handles both the case of integer and floating-point exponents. Moreover, when the exponent is an array, the function will calculate the element-wise powers of the input values.

Parameters:
  • x (int | float | drjit.ArrayBase) – A Python or Dr.Jit array type as input value

  • y (int | float | drjit.ArrayBase) – A Python or Dr.Jit array type as exponent

Returns:

input value raised to the power

Return type:

int | float | drjit.ArrayBase

Accuracy (single precision)#

Note

The trigonometric functions sin, cos, and tan are optimized for low error on the domain \(|x| < 8192\) and don’t perform as well beyond this range.

Function

Tested domain

Abs. error (mean)

Abs. error (max)

Rel. error (mean)

Rel. error (max)

\(\text{sin}()\)

\(-8192 < x < 8192\)

\(1.2 \cdot 10^{-8}\)

\(1.2 \cdot 10^{-7}\)

\(1.9 \cdot 10^{-8}\,(0.25\,\text{ulp})\)

\(1.8 \cdot 10^{-6}\,(19\,\text{ulp})\)

\(\text{cos}()\)

\(-8192 < x < 8192\)

\(1.2 \cdot 10^{-8}\)

\(1.2 \cdot 10^{-7}\)

\(1.9 \cdot 10^{-8}\,(0.25\,\text{ulp})\)

\(3.1 \cdot 10^{-6}\,(47\,\text{ulp})\)

\(\text{tan}()\)

\(-8192 < x < 8192\)

\(4.7 \cdot 10^{-6}\)

\(8.1 \cdot 10^{-1}\)

\(3.4 \cdot 10^{-8}\,(0.42\,\text{ulp})\)

\(3.1 \cdot 10^{-6}\,(30\,\text{ulp})\)

\(\text{asin}()\)

\(-1 < x < 1\)

\(2.3 \cdot 10^{-8}\)

\(1.2 \cdot 10^{-7}\)

\(2.9 \cdot 10^{-8}\,(0.33\,\text{ulp})\)

\(2.3 \cdot 10^{-7}\,(2\,\text{ulp})\)

\(\text{acos}()\)

\(-1 < x < 1\)

\(4.7 \cdot 10^{-8}\)

\(2.4 \cdot 10^{-7}\)

\(2.9 \cdot 10^{-8}\,(0.33\,\text{ulp})\)

\(1.2 \cdot 10^{-7}\,(1\,\text{ulp})\)

\(\text{atan}()\)

\(-1 < x < 1\)

\(1.8 \cdot 10^{-7}\)

\(6 \cdot 10^{-7}\)

\(4.2 \cdot 10^{-7}\,(4.9\,\text{ulp})\)

\(8.2 \cdot 10^{-7}\,(12\,\text{ulp})\)

\(\text{sinh}()\)

\(-10 < x < 10\)

\(2.6 \cdot 10^{-5}\)

\(2 \cdot 10^{-3}\)

\(2.8 \cdot 10^{-8}\,(0.34\,\text{ulp})\)

\(2.7 \cdot 10^{-7}\,(3\,\text{ulp})\)

\(\text{cosh}()\)

\(-10 < x < 10\)

\(2.9 \cdot 10^{-5}\)

\(2 \cdot 10^{-3}\)

\(2.9 \cdot 10^{-8}\,(0.35\,\text{ulp})\)

\(2.5 \cdot 10^{-7}\,(4\,\text{ulp})\)

\(\text{tanh}()\)

\(-10 < x < 10\)

\(4.8 \cdot 10^{-8}\)

\(4.2 \cdot 10^{-7}\)

\(5 \cdot 10^{-8}\,(0.76\,\text{ulp})\)

\(5 \cdot 10^{-7}\,(7\,\text{ulp})\)

\(\text{asinh}()\)

\(-30 < x < 30\)

\(2.8 \cdot 10^{-8}\)

\(4.8 \cdot 10^{-7}\)

\(1 \cdot 10^{-8}\,(0.13\,\text{ulp})\)

\(1.7 \cdot 10^{-7}\,(2\,\text{ulp})\)

\(\text{acosh}()\)

\(1 < x < 10\)

\(2.9 \cdot 10^{-8}\)

\(2.4 \cdot 10^{-7}\)

\(1.5 \cdot 10^{-8}\,(0.18\,\text{ulp})\)

\(2.4 \cdot 10^{-7}\,(3\,\text{ulp})\)

\(\text{atanh}()\)

\(-1 < x < 1\)

\(9.9 \cdot 10^{-9}\)

\(2.4 \cdot 10^{-7}\)

\(1.5 \cdot 10^{-8}\,(0.18\,\text{ulp})\)

\(1.2 \cdot 10^{-7}\,(1\,\text{ulp})\)

\(\text{exp}()\)

\(-20 < x < 30\)

\(0.72 \cdot 10^{4}\)

\(0.1 \cdot 10^{7}\)

\(2.4 \cdot 10^{-8}\,(0.27\,\text{ulp})\)

\(1.2 \cdot 10^{-7}\,(1\,\text{ulp})\)

\(\text{log}()\)

\(10^{-20} < x < 2\cdot 10^{30}\)

\(9.6 \cdot 10^{-9}\)

\(7.6 \cdot 10^{-6}\)

\(1.4 \cdot 10^{-10}\,(0.0013\,\text{ulp})\)

\(1.2 \cdot 10^{-7}\,(1\,\text{ulp})\)

\(\text{erf}()\)

\(-1 < x < 1\)

\(3.2 \cdot 10^{-8}\)

\(1.8 \cdot 10^{-7}\)

\(6.4 \cdot 10^{-8}\,(0.78\,\text{ulp})\)

\(3.3 \cdot 10^{-7}\,(4\,\text{ulp})\)

\(\text{erfc}()\)

\(-1 < x < 1\)

\(3.4 \cdot 10^{-8}\)

\(2.4 \cdot 10^{-7}\)

\(6.4 \cdot 10^{-8}\,(0.79\,\text{ulp})\)

\(1 \cdot 10^{-6}\,(11\,\text{ulp})\)

Accuracy (double precision)#

Function

Tested domain

Abs. error (mean)

Abs. error (max)

Rel. error (mean)

Rel. error (max)

\(\text{sin}()\)

\(-8192 < x < 8192\)

\(2.2 \cdot 10^{-17}\)

\(2.2 \cdot 10^{-16}\)

\(3.6 \cdot 10^{-17}\,(0.25\,\text{ulp})\)

\(3.1 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{cos}()\)

\(-8192 < x < 8192\)

\(2.2 \cdot 10^{-17}\)

\(2.2 \cdot 10^{-16}\)

\(3.6 \cdot 10^{-17}\,(0.25\,\text{ulp})\)

\(3 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{tan}()\)

\(-8192 < x < 8192\)

\(6.8 \cdot 10^{-16}\)

\(1.2 \cdot 10^{-10}\)

\(5.4 \cdot 10^{-17}\,(0.35\,\text{ulp})\)

\(4.1 \cdot 10^{-16}\,(3\,\text{ulp})\)

\(\text{cot}()\)

\(-8192 < x < 8192\)

\(4.9 \cdot 10^{-16}\)

\(1.2 \cdot 10^{-10}\)

\(5.5 \cdot 10^{-17}\,(0.36\,\text{ulp})\)

\(4.4 \cdot 10^{-16}\,(3\,\text{ulp})\)

\(\text{asin}()\)

\(-1 < x < 1\)

\(1.3 \cdot 10^{-17}\)

\(2.2 \cdot 10^{-16}\)

\(1.5 \cdot 10^{-17}\,(0.098\,\text{ulp})\)

\(2.2 \cdot 10^{-16}\,(1\,\text{ulp})\)

\(\text{acos}()\)

\(-1 < x < 1\)

\(5.4 \cdot 10^{-17}\)

\(4.4 \cdot 10^{-16}\)

\(3.5 \cdot 10^{-17}\,(0.23\,\text{ulp})\)

\(2.2 \cdot 10^{-16}\,(1\,\text{ulp})\)

\(\text{atan}()\)

\(-1 < x < 1\)

\(4.3 \cdot 10^{-17}\)

\(3.3 \cdot 10^{-16}\)

\(1 \cdot 10^{-16}\,(0.65\,\text{ulp})\)

\(7.1 \cdot 10^{-16}\,(5\,\text{ulp})\)

\(\text{sinh}()\)

\(-10 < x < 10\)

\(3.1 \cdot 10^{-14}\)

\(1.8 \cdot 10^{-12}\)

\(3.3 \cdot 10^{-17}\,(0.22\,\text{ulp})\)

\(4.3 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{cosh}()\)

\(-10 < x < 10\)

\(2.2 \cdot 10^{-14}\)

\(1.8 \cdot 10^{-12}\)

\(2 \cdot 10^{-17}\,(0.13\,\text{ulp})\)

\(2.9 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{tanh}()\)

\(-10 < x < 10\)

\(5.6 \cdot 10^{-17}\)

\(3.3 \cdot 10^{-16}\)

\(6.1 \cdot 10^{-17}\,(0.52\,\text{ulp})\)

\(5.5 \cdot 10^{-16}\,(3\,\text{ulp})\)

\(\text{asinh}()\)

\(-30 < x < 30\)

\(5.1 \cdot 10^{-17}\)

\(8.9 \cdot 10^{-16}\)

\(1.9 \cdot 10^{-17}\,(0.13\,\text{ulp})\)

\(4.4 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{acosh}()\)

\(1 < x < 10\)

\(4.9 \cdot 10^{-17}\)

\(4.4 \cdot 10^{-16}\)

\(2.6 \cdot 10^{-17}\,(0.17\,\text{ulp})\)

\(6.6 \cdot 10^{-16}\,(5\,\text{ulp})\)

\(\text{atanh}()\)

\(-1 < x < 1\)

\(1.8 \cdot 10^{-17}\)

\(4.4 \cdot 10^{-16}\)

\(3.2 \cdot 10^{-17}\,(0.21\,\text{ulp})\)

\(3 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{exp}()\)

\(-20 < x < 30\)

\(4.7 \cdot 10^{-6}\)

\(2 \cdot 10^{-3}\)

\(2.5 \cdot 10^{-17}\,(0.16\,\text{ulp})\)

\(3.3 \cdot 10^{-16}\,(2\,\text{ulp})\)

\(\text{log}()\)

\(10^{-20} < x < 2\cdot 10^{30}\)

\(1.9 \cdot 10^{-17}\)

\(1.4 \cdot 10^{-14}\)

\(2.7 \cdot 10^{-19}\,(0.0013\,\text{ulp})\)

\(2.2 \cdot 10^{-16}\,(1\,\text{ulp})\)

\(\text{erf}()\)

\(-1 < x < 1\)

\(4.7 \cdot 10^{-17}\)

\(4.4 \cdot 10^{-16}\)

\(9.6 \cdot 10^{-17}\,(0.63\,\text{ulp})\)

\(5.9 \cdot 10^{-16}\,(5\,\text{ulp})\)

\(\text{erfc}()\)

\(-1 < x < 1\)

\(4.8 \cdot 10^{-17}\)

\(4.4 \cdot 10^{-16}\)

\(9.6 \cdot 10^{-17}\,(0.64\,\text{ulp})\)

\(2.5 \cdot 10^{-15}\,(16\,\text{ulp})\)

Safe mathematical functions#

drjit.safe_sqrt(arg)#

Safely evaluate the square root of the provided input avoiding domain errors.

Negative inputs produce a 0.0 output value.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Square root of the input

Return type:

float | drjit.ArrayBase

drjit.safe_asin(arg)#

Safe wrapper around drjit.asin() that avoids domain errors.

Input values are clipped to the \((-1, 1)\) domain.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Arcsine approximation

Return type:

float | drjit.ArrayBase

drjit.safe_acos(arg)#

Safe wrapper around drjit.acos() that avoids domain errors.

Input values are clipped to the \((-1, 1)\) domain.

Parameters:

arg (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type

Returns:

Arccosine approximation

Return type:

float | drjit.ArrayBase

Constants#

drjit.epsilon(t)#

Returns the machine epsilon.

The machine epsilon gives an upper bound on the relative approximation error due to rounding in floating point arithmetic.

Parameters:

t (type) – Python or Dr.Jit type determining whether to consider 32 or 64 bits floating point precision.

Returns:

machine epsilon

Return type:

float

drjit.one_minus_epsilon(t)#

Returns one minus machine epsilon value.

Parameters:

t (type) – Python or Dr.Jit type determining whether to consider 32 or 64 bits floating point precision.

Returns:

one minus machine epsilon value

Return type:

float

drjit.recip_overflow(t)#

Returns the reciprocal overflow threshold value.

Any numbers below this threshold will overflow to infinity when a reciprocal is evaluated.

Parameters:

t (type) – Python or Dr.Jit type determining whether to consider 32 or 64 bits floating point precision.

Returns:

reciprocal overflow threshold value

Return type:

float

drjit.smallest(t)#

Returns the smallest normalized floating point value.

Parameters:

t (type) – Python or Dr.Jit type determining whether to consider 32 or 64 bits floating point precision.

Returns:

smallest normalized floating point value

Return type:

float

drjit.largest(t)#

Returns the largest normalized floating point value.

Parameters:

t (type) – Python or Dr.Jit type determining whether to consider 32 or 64 bits floating point precision.

Returns:

largest normalized floating point value

Return type:

float

Array base class#

class drjit.ArrayBase#
__len__()#
__iter__()#
__repr__()#

Return repr(self).

__bool__()#

Casts the array to a Python bool type. This is only permissible when self represents an boolean array of both depth and size 1.

__add__(b)#
__radd__(b)#
__iadd__(b)#
__sub__(b)#
__rsub__(b)#
__isub__(b)#
__mul__(b)#
__rmul__(b)#
__imul__(b)#
__truediv__(b)#
__rtruediv__(b)#
__itruediv__(b)#
__floordiv__(b)#
__rfloordiv__(b)#
__ifloordiv__(b)#
__mod__(b)#
__imod__(b)#
__rshift__(b)#
__rrshift__(b)#
__irshift__(b)#
__lshift__(b)#
__rlshift__(b)#
__ilshift__(b)#
__and__(b)#
__rand__(b)#
__iand__(b)#
__or__(b)#

Return self|value.

__ror__(b)#

Return value|self.

__ior__(b)#
__xor__(b)#
__rxor__(b)#
__ixor__(b)#
__abs__()#
__le__(b)#

Return self<=value.

__lt__(b)#

Return self<value.

__ge__(b)#

Return self>=value.

__gt__(b)#

Return self>value.

__ne__(value, /)#

Return self!=value.

__eq__(b)#

Return self==value.

__dlpack__()#

Automatic differentiation#

drjit.detach(arg, preserve_type=True)#

Transforms the input variable into its non-differentiable version (detaches it from the AD computational graph).

This function is able to traverse data-structures such a sequences, mappings or custom data structure and applies the transformation to the underlying variables.

When the input variable isn’t a Dr.Jit differentiable array, it is returned as it is.

While the type of the returned array is preserved by default, it is possible to set the preserve_type argument to false to force the returned type to be non-differentiable.

Parameters:
  • arg (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequence, or mapping.

  • preserve_type (bool) – Defines whether the returned variable should preserve the type of the input variable.

Returns:

The detached variable.

Return type:

object

drjit.enable_grad(*args)#

Enable gradient tracking for the provided variables.

This function accepts a variable-length list of arguments and processes it as follows:

  • It recurses into sequences (tuple, list, etc.)

  • It recurses into the values of mappings (dict, etc.)

  • It recurses into the fields of custom data structures.

During recursion, the function enables gradient tracking for all Dr.Jit arrays. For every other types, this function won’t do anything.

Parameters:

*args (tuple) – A variable-length list of Dr.Jit array instances, custom data structures, sequences, or mappings.

drjit.disable_grad(*args)#

Disable gradient tracking for the provided variables.

This function accepts a variable-length list of arguments and processes it as follows:

  • It recurses into sequences (tuple, list, etc.)

  • It recurses into the values of mappings (dict, etc.)

  • It recurses into the fields of custom data structures.

During recursion, the function disables gradient tracking for all Dr.Jit arrays. For every other types, this function won’t do anything.

Parameters:

*args (tuple) – A variable-length list of Dr.Jit array instances, custom data structures, sequences, or mappings.

drjit.set_grad_enabled(arg, value)#

Enable or disable gradient tracking on the provided variables.

Parameters:
  • arg (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequence, or mapping.

  • value (bool) – Defines whether gradient tracking should be enabled or disabled.

drjit.grad_enabled(*args)#

Return whether gradient tracking is enabled on any of the given variables.

Parameters:

*args (tuple) – A variable-length list of Dr.Jit array instances, custom data structures, sequences, or mappings. The function will recursively traverse data structures to discover all Dr.Jit arrays.

Returns:

True if any variable has gradient tracking enabled, False otherwise.

Return type:

bool

drjit.grad(arg, preserve_type=True)#

Return the gradient value associated to a given variable.

When the variable doesn’t have gradient tracking enabled, this function returns 0.

For all input variables that are not Dr.Jit arrays or mapping and sequences, thi function returns None.

Parameters:
  • arg (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequences, or mapping.

  • preserve_type (bool) – Defines whether the returned variable should preserve the type of the input variable.

Returns:

the gradient value associated to the input variable.

Return type:

object

drjit.set_grad(dst, src)#

Set the gradient value to the provided variable.

Broadcasting is applied to the gradient value if necessary and possible to match the type of the input variable.

Parameters:
  • dst (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequences, or mapping.

  • src (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequences, or mapping.

drjit.accum_grad(dst, src)#

Accumulate into the gradient of a variable.

Broadcasting is applied to the gradient value if necessary and possible to match the type of the input variable.

Parameters:
  • dst (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequences, or mapping.

  • src (object) – An arbitrary Dr.Jit array, tensor, custom data structure, sequences, or mapping.

drjit.replace_grad(dst, src)#

Replace the gradient value of dst with the one of src.

Broadcasting is applied to dst if necessary to match the type of src.

Parameters:
  • dst (object) – An arbitrary Dr.Jit array, tensor, or scalar builtin instance.

  • src (object) – An differentiable Dr.Jit array or tensor.

Returns:

the variable with the replaced gradients.

Return type:

object

drjit.traverse(dtype, mode, flags=<ADFlag.Default: 7>)#

Propagate derivatives through the enqueued set of edges in the AD computational graph in the direction specified by mode.

By default, Dr.Jit’s AD system destructs the enqueued input graph during AD traversal. This frees up resources, which is useful when working with large wavefronts or very complex computation graphs. However, this also prevents repeated propagation of gradients through a shared subgraph that is being differentiated multiple times.

To support more fine-grained use cases that require this, the following flags can be used to control what should and should not be destructed:

  • ADFlag.ClearNone: clear nothing

  • ADFlag.ClearEdges: delete all traversed edges from the computation graph

  • ADFlag.ClearInput: clear the gradients of processed input vertices (in-degree == 0)

  • ADFlag.ClearInterior: clear the gradients of processed interior vertices (out-degree != 0)

  • ADFlag.ClearVertices: clear gradients of processed vertices only, but leave edges intact

  • ADFlag.Default: clear everything (default behaviour)

Parameters:
  • dtype (type) – defines the Dr.JIT array type used to build the AD graph

  • mode (ADMode) – defines the mode traversal (backward or forward)

  • flags (ADFlag | int) – flags to control what should and should not be destructed during forward/backward mode traversal.

drjit.enqueue(mode, *args)#

Enqueues variable for the subsequent AD traversal.

In Dr.Jit, the process of automatic differentiation is split into two parts:

  1. Discover and enqueue the variables to be considered as inputs during the subsequent AD traversal.

  2. Traverse the AD graph starting from the enqueued variables to propagate the gradients towards the output variables (e.g. leaf in the AD graph).

This function handles the first part can operate in different modes depending on the specified mode:

  • ADMode.Forward: the provided value will be considered as input during the subsequent AD traversal.

  • ADMode.Backward: a traversal of the AD graph starting from the provided value will take place to find all potential source of gradients and enqueue them.

For example, a typical chain of operations to forward propagate the gradients from a to b would look as follow:

a = dr.llvm.ad.Float(1.0)
dr.enable_grad(a)
b = f(a) # some computation involving `a`
dr.set_gradient(a, 1.0)
dr.enqueue(dr.ADMode.Forward, a)
dr.traverse(dr.llvm.ad.Float, dr.ADMode.Forward)
grad = dr.grad(b)

It could be the case that f(a) involves other differentiable variables that already contain some gradients. In this situation we can use ADMode.Backward to discover and enqueue them before the traversal.

a = dr.llvm.ad.Float(1.0)
dr.enable_grad(a)
b = f(a, ...) # some computation involving `a` and some hidden variables
dr.set_gradient(a, 1.0)
dr.enqueue(dr.ADMode.Backward, b)
dr.traverse(dr.llvm.ad.Float, dr.ADMode.Forward)
grad = dr.grad(b)

Dr.Jit also provides a higher level API that encapsulate this logic in a few different functions:

Parameters:
  • mode (ADMode) – defines the enqueuing mode (backward or forward)

  • *args (tuple) – A variable-length list of Dr.Jit array instances, tensors, custom data structures, sequences, or mappings.

drjit.forward_from(arg, flags=<ADFlag.Default: 7>)#

Forward propagates gradients from a provided Dr.Jit differentiable array.

This function will first see the gradient value of the provided variable to 1.0 before executing the AD graph traversal.

An exception will be raised when the provided array doesn’t have gradient tracking enabled or if it isn’t an instance of a Dr.Jit differentiable array type.

Parameters:
  • arg (object) – A Dr.Jit differentiable array instance.

  • flags (ADFlag | int) – flags to control what should and should not be destructed during the traversal. The default value is ADFlag.Default.

drjit.forward(arg, flags=<ADFlag.Default: 7>)#

Forward propagates gradients from a provided Dr.Jit differentiable array.

This function will first see the gradient value of the provided variable to 1.0 before executing the AD graph traversal.

An exception will be raised when the provided array doesn’t have gradient tracking enabled or if it isn’t an instance of a Dr.Jit differentiable array type.

This function is an alias of drjit.forward_from().

Parameters:
  • arg (object) – A Dr.Jit differentiable array instance.

  • flags (ADFlag | int) – flags to control what should and should not be destructed during the traversal. The default value is ADFlag.Default.

drjit.forward_to(*args, flags=<ADFlag.Default: 7>)#

Forward propagates gradients to a set of provided Dr.Jit differentiable arrays.

Internally, the AD computational graph will be first traversed backward to find all potential source of gradient for the provided array. Then only the forward gradient propagation traversal takes place.

The flags argument should be provided as a keyword argument for this function.

An exception will be raised when the provided array doesn’t have gradient tracking enabled or if it isn’t an instance of a Dr.Jit differentiable array type.

Parameters:
  • *args (tuple) – A variable-length list of Dr.Jit differentiable array, tensor, custom data structure, sequences, or mapping.

  • flags (ADFlag | int) – flags to control what should and should not be destructed during the traversal. The default value is ADFlag.Default.

Returns:

the gradient value associated to the output variables.

Return type:

object

drjit.backward_from(arg, flags=<ADFlag.Default: 7>)#

Backward propagates gradients from a provided Dr.Jit differentiable array.

An exception will be raised when the provided array doesn’t have gradient tracking enabled or if it isn’t an instance of a Dr.Jit differentiable array type.

Parameters:
  • arg (object) – A Dr.Jit differentiable array instance.

  • flags (ADFlag | int) – flags to control what should and should not be destructed during the traversal. The default value is ADFlag.Default.

drjit.backward(arg, flags=<ADFlag.Default: 7>)#

Backward propagate gradients from a provided Dr.Jit differentiable array.

An exception will be raised when the provided array doesn’t have gradient tracking enabled or if it isn’t an instance of a Dr.Jit differentiable array type.

This function is an alias of drjit.backward_from().

Parameters:
  • arg (object) – A Dr.Jit differentiable array instance.

  • flags (ADFlag | int) – flags to control what should and should not be destructed during the traversal. The default value is ADFlag.Default.

drjit.backward_to(*args, flags=<ADFlag.Default: 7>)#

Backward propagate gradients to a set of provided Dr.Jit differentiable arrays.

Internally, the AD computational graph will be first traversed forward to find all potential source of gradient for the provided array. Then only the backward gradient propagation traversal takes place.

The flags argument should be provided as a keyword argument for this function.

An exception will be raised when the provided array doesn’t have gradient tracking enabled or if it isn’t an instance of a Dr.Jit differentiable array type.

Parameters:
  • *args (tuple) – A variable-length list of Dr.Jit differentiable array, tensor, custom data structure, sequences, or mapping.

  • flags (ADFlag | int) – flags to control what should and should not be destructed during the traversal. The default value is ADFlag.Default.

Returns:

the gradient value associated to the output variables.

Return type:

object

drjit.suspend_grad(*args, when=True)#

Context manager for temporally suspending derivative tracking.

Dr.Jit’s AD layer keeps track of a set of variables for which derivative tracking is currently enabled. Using this context manager is it possible to define a scope in which variables will be subtracted from that set, thereby controlling what derivative terms shouldn’t be generated in that scope.

The variables to be subtracted from the current set of enabled variables can be provided as function arguments. If none are provided, the scope defined by this context manager will temporally disable all derivative tracking.

a = dr.llvm.ad.Float(1.0)
b = dr.llvm.ad.Float(2.0)
dr.enable_grad(a, b)

with suspend_grad(): # suspend all derivative tracking
    c = a + b

assert not dr.grad_enabled(c)

with suspend_grad(a): # only suspend derivative tracking on `a`
    d = 2.0 * a
    e = 4.0 * b

assert not dr.grad_enabled(d)
assert dr.grad_enabled(e)

In a scope where derivative tracking is completely suspended, the AD layer will ignore any attempt to enable gradient tracking on a variable:

a = dr.llvm.ad.Float(1.0)

with suspend_grad():
    dr.enable_grad(a) # <-- ignored
    assert not dr.grad_enabled(a)

assert not dr.grad_enabled(a)

The optional when boolean keyword argument can be defined to specifed a condition determining whether to suspend the tracking of derivatives or not.

a = dr.llvm.ad.Float(1.0)
dr.enable_grad(a)

cond = condition()

with suspend_grad(when=cond):
    b = 4.0 * a

assert dr.grad_enabled(b) == not cond
Parameters:
  • *args (tuple) – A variable-length list of differentiable Dr.Jit array instances, custom data structures, sequences, or mappings. The function will recursively traverse data structures to discover all Dr.Jit arrays.

  • when (bool) – An optional Python boolean determining whether to suspend derivative tracking.

drjit.resume_grad(*args, when=True)#

Context manager for temporally resume derivative tracking.

Dr.Jit’s AD layer keeps track of a set of variables for which derivative tracking is currently enabled. Using this context manager is it possible to define a scope in which variables will be added to that set, thereby controlling what derivative terms should be generated in that scope.

The variables to be added to the current set of enabled variables can be provided as function arguments. If none are provided, the scope defined by this context manager will temporally resume derivative tracking for all variables.

a = dr.llvm.ad.Float(1.0)
b = dr.llvm.ad.Float(2.0)
dr.enable_grad(a, b)

with suspend_grad():
    c = a + b

    with resume_grad():
        d = a + b

    with resume_grad(a):
        e = 2.0 * a
        f = 4.0 * b

assert not dr.grad_enabled(c)
assert dr.grad_enabled(d)
assert dr.grad_enabled(e)
assert not dr.grad_enabled(f)

The optional when boolean keyword argument can be defined to specifed a condition determining whether to resume the tracking of derivatives or not.

a = dr.llvm.ad.Float(1.0)
dr.enable_grad(a)

cond = condition()

with suspend_grad():
    with resume_grad(when=cond):
        b = 4.0 * a

assert dr.grad_enabled(b) == cond
Parameters:
  • *args (tuple) – A variable-length list of differentiable Dr.Jit array instances, custom data structures, sequences, or mappings. The function will recursively traverse data structures to discover all Dr.Jit arrays.

  • when (bool) – An optional Python boolean determining whether to resume derivative tracking.

drjit.isolate_grad(when=True)#

Context manager to temporarily isolate outside world from AD traversals.

Dr.Jit provides isolation boundaries to postpone AD traversals steps leaving a specific scope. For instance this function is used internally to implement differentiable loops and polymorphic calls.

drjit.graphviz_ad(as_str=False)#

Assembles a graphviz diagram for the computational graph trace by the AD system.

Parameters:

as_str (bool) – whether the function should return the graphviz object as a string representation or not.

Returns:

the graphviz obj (or its string representation).

Return type:

object

class drjit.CustomOp#

Base class to implement custom differentiable operations.

Dr.Jit can compute derivatives of builtin operations in both forward and reverse mode. In some cases, it may be useful or even necessary to tell Dr.Jit how a particular operation should be differentiated.

This can be achieved by extending this class, overwriting callback functions that will later be invoked when the AD backend traverses the associated node in the computation graph. This class also provides a convenient way of stashing temporary results during the original function evaluation that can be accessed later on as part of forward or reverse-mode differentiation.

Look at the section on AD custom operations for more detailed information.

A class that inherits from this class should override a few methods as done in the code snippet below. dr.custom() can then be used to evaluate the custom operation and properly attach it to the AD graph.

class MyCustomOp(dr.CustomOp):
    def eval(self, *args):
        # .. evaluate operation ..

    def forward(self):
        # .. compute forward-mode derivatives ..

    def backward(self):
        # .. compute backward-mode derivatives ..

    def name(self):
        return "MyCustomOp[]"

dr.custom(MyCustomOp, *args)
eval(self, *args) object#

Evaluate the custom function in primal mode.

The inputs will be detached from the AD graph, and the output must also be detached.

Danger

This method must be overriden, no default implementation provided.

forward()#

Evaluated forward-mode derivatives.

Danger

This method must be overriden, no default implementation provided.

backward()#

Evaluated backward-mode derivatives.

Danger

This method must be overriden, no default implementation provided.

name()#

Return a descriptive name of the CustomOp instance.

The name returned by this method is used in the GraphViz output.

If not overriden, this method returns "CustomOp[unnamed]".

grad_out()#

Access the gradient associated with the output argument (backward mode AD).

Returns:

the gradient value associated with the output argument.

Return type:

object

set_grad_out(value)#

Accumulate a gradient value into the output argument (forward mode AD).

Parameters:

value (object) – gradient value to accumulate.

grad_in(name)#

Access the gradient associated with the input argument name (fwd. mode AD).

Parameters:

name (str) – name associated to an input variable (e.g. keyword argument).

Returns:

the gradient value associated with the input argument.

Return type:

object

set_grad_in(name, value)#

Accumulate a gradient value into an input argument (backward mode AD).

Parameters:
  • name (str) – name associated to the input variable (e.g. keyword argument).

  • value (object) – gradient value to accumulate.

add_input(value)#

Register an implicit input dependency of the operation on an AD variable.

This function should be called by the eval() implementation when an operation has a differentiable dependence on an input that is not an input argument (e.g. a private instance variable).

Parameters:

value (object) – variable this operation depends on implicitly.

add_output(value)#

Register an implicit output dependency of the operation on an AD variable.

This function should be called by the

ef eval() implementation when an

operation has a differentiable dependence on an output that is not an return value of the operation (e.g. a private instance variable).

Args:

value (object): variable this operation depends on implicitly.

drjit.custom(cls, *args, **kwargs)#

Evaluate a custom differentiable operation (see CustomOp).

Look at the section on AD custom operations for more detailed information.

drjit.wrap_ad(source: str, target: str)#

Function decorator that wraps the excecution of a function using a different AD framework to ensure that gradients can flow seamlessly between both frameworks.

Using this decorator it is possible to mix AD-aware computation between two AD frameworks (e.g. Dr.Jit and PyTorch). The wrapped function’s arguments will be casted to the tensor types corresponding to the target framework. Similarily, the return values will be casted to the tensor types corresponding to the source framework.

The decorated function can take an arbitrary number of arguments and have any number of return values.

Currently only the following combination of frameworks are supported:

source

target

Forward AD

Backward AD

drjit

torch

✔️

torch

drjit

✔️

The example below shows how to wrap a Dr.Jit function in a PyTorch script:

# Start with a PyTorch tensor
a = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Some PyTorch arithmetic
b = torch.sin(a)

# Wrap a function performing some arithmetic using Dr.Jit
@dr.wrap_ad(source='torch', target='drjit')
def dr_func(x):
    return dr.cos(x) + dr.power(x, 2)

# Excecute the wrapped function (returns a PyTorch tensor)
c = dr_func(b)

# Some more PyTorch arithmetic
d = torch.tan(c)

# Propagate gradients to variable a (through PyTorch -> Dr.Jit -> PyTorch)
d.sum().backward()

# Inspect the resulting gradients
print(a.grad)

Similarily the following example shows how to wrap PyTorch code into a Dr.Jit script:

# Start with a Dr.Jit tensor
a = dr.llvm.ad.TensorXf([1, 2, 3], shape=[3])
dr.enable_grad(a)

# Some Dr.Jit arithmetic
b = dr.sin(a)

# Wrap a function performing some arithmetic using PyTorch
@dr.wrap_ad(source='drjit', target='torch')
def torch_func(x):
    return torch.cos(x) + torch.sin(x)

# Excecute the wrapped function (returns a Dr.Jit tensor)
c = torch_func(b)

# Some more Dr.Jit arithmetic
d = dr.tan(c)

# Propagate gradients to variable a (through Dr.Jit -> PyTorch -> Dr.Jit)
dr.backward(d)

# Inspect the resulting gradients
print(dr.grad(a))

Danger

Forward-mode AD isn’t currently supported by this operation.

Parameters:
  • source (str | module) – The AD framework used outside of the wrapped function.

  • target (str | module) – The AD framework used within the wrapped function.

Returns:

The decorated function.

Concrete array classes#

Scalar array namespace (drjit.scalar)#

drjit.Bool: type = bool#
drjit.Float: type = float#
drjit.Float64: type = float#
drjit.Int: type = int#
drjit.Int64: type = int#
drjit.UInt: type = int#
drjit.UInt64: type = int#
class drjit.scalar.Array0b#

Bases: ArrayBase

class drjit.scalar.Array1b#

Bases: ArrayBase

class drjit.scalar.Array2b#

Bases: ArrayBase

class drjit.scalar.Array3b#

Bases: ArrayBase

class drjit.scalar.Array4b#

Bases: ArrayBase

class drjit.scalar.ArrayXb#

Bases: ArrayBase

class drjit.scalar.Array0f#

Bases: ArrayBase

class drjit.scalar.Array1f#

Bases: ArrayBase

class drjit.scalar.Array2f#

Bases: ArrayBase

class drjit.scalar.Array3f#

Bases: ArrayBase

class drjit.scalar.Array4f#

Bases: ArrayBase

class drjit.scalar.ArrayXf#

Bases: ArrayBase

class drjit.scalar.Array0u#

Bases: ArrayBase

class drjit.scalar.Array1u#

Bases: ArrayBase

class drjit.scalar.Array2u#

Bases: ArrayBase

class drjit.scalar.Array3u#

Bases: ArrayBase

class drjit.scalar.Array4u#

Bases: ArrayBase

class drjit.scalar.ArrayXu#

Bases: ArrayBase

class drjit.scalar.Array0i#

Bases: ArrayBase

class drjit.scalar.Array1i#

Bases: ArrayBase

class drjit.scalar.Array2i#

Bases: ArrayBase

class drjit.scalar.Array3i#

Bases: ArrayBase

class drjit.scalar.Array4i#

Bases: ArrayBase

class drjit.scalar.ArrayXi#

Bases: ArrayBase

class drjit.scalar.Array0f64#

Bases: ArrayBase

class drjit.scalar.Array1f64#

Bases: ArrayBase

class drjit.scalar.Array2f64#

Bases: ArrayBase

class drjit.scalar.Array3f64#

Bases: ArrayBase

class drjit.scalar.Array4f64#

Bases: ArrayBase

class drjit.scalar.ArrayXf64#

Bases: ArrayBase

class drjit.scalar.Array0u64#

Bases: ArrayBase

class drjit.scalar.Array1u64#

Bases: ArrayBase

class drjit.scalar.Array2u64#

Bases: ArrayBase

class drjit.scalar.Array3u64#

Bases: ArrayBase

class drjit.scalar.Array4u64#

Bases: ArrayBase

class drjit.scalar.ArrayXu64#

Bases: ArrayBase

class drjit.scalar.Array0i64#

Bases: ArrayBase

class drjit.scalar.Array1i64#

Bases: ArrayBase

class drjit.scalar.Array2i64#

Bases: ArrayBase

class drjit.scalar.Array3i64#

Bases: ArrayBase

class drjit.scalar.Array4i64#

Bases: ArrayBase

class drjit.scalar.ArrayXi64#

Bases: ArrayBase

class drjit.scalar.Matrix2u#

Bases: ArrayBase

class drjit.scalar.Matrix2i#

Bases: ArrayBase

class drjit.scalar.Matrix2f#

Bases: ArrayBase

class drjit.scalar.Matrix2f64#

Bases: ArrayBase

class drjit.scalar.Matrix3u#

Bases: ArrayBase

class drjit.scalar.Matrix3i#

Bases: ArrayBase

class drjit.scalar.Matrix3f#

Bases: ArrayBase

class drjit.scalar.Matrix3f64#

Bases: ArrayBase

class drjit.scalar.Matrix4u#

Bases: ArrayBase

class drjit.scalar.Matrix4i#

Bases: ArrayBase

class drjit.scalar.Matrix4f#

Bases: ArrayBase

class drjit.scalar.Matrix4f64#

Bases: ArrayBase

class drjit.scalar.Complex2f#

Bases: ArrayBase

class drjit.scalar.Complex2f64#

Bases: ArrayBase

class drjit.scalar.Quaternion4f#

Bases: ArrayBase

class drjit.scalar.Quaternion4f64#

Bases: ArrayBase

CUDA array namespace (drjit.cuda)#

class drjit.cuda.Bool#

Bases: ArrayBase

class drjit.cuda.Float#

Bases: ArrayBase

class drjit.cuda.Float64#

Bases: ArrayBase

class drjit.cuda.UInt#

Bases: ArrayBase

class drjit.cuda.UInt64#

Bases: ArrayBase

class drjit.cuda.Int#

Bases: ArrayBase

class drjit.cuda.Int64#

Bases: ArrayBase

class drjit.cuda.Array0b#

Bases: ArrayBase

class drjit.cuda.Array1b#

Bases: ArrayBase

class drjit.cuda.Array2b#

Bases: ArrayBase

class drjit.cuda.Array3b#

Bases: ArrayBase

class drjit.cuda.Array4b#

Bases: ArrayBase

class drjit.cuda.ArrayXb#

Bases: ArrayBase

class drjit.cuda.Array0f#

Bases: ArrayBase

class drjit.cuda.Array1f#

Bases: ArrayBase

class drjit.cuda.Array2f#

Bases: ArrayBase

class drjit.cuda.Array3f#

Bases: ArrayBase

class drjit.cuda.Array4f#

Bases: ArrayBase

class drjit.cuda.ArrayXf#

Bases: ArrayBase

class drjit.cuda.Array0u#

Bases: ArrayBase

class drjit.cuda.Array1u#

Bases: ArrayBase

class drjit.cuda.Array2u#

Bases: ArrayBase

class drjit.cuda.Array3u#

Bases: ArrayBase

class drjit.cuda.Array4u#

Bases: ArrayBase

class drjit.cuda.ArrayXu#

Bases: ArrayBase

class drjit.cuda.Array0i#

Bases: ArrayBase

class drjit.cuda.Array1i#

Bases: ArrayBase

class drjit.cuda.Array2i#

Bases: ArrayBase

class drjit.cuda.Array3i#

Bases: ArrayBase

class drjit.cuda.Array4i#

Bases: ArrayBase

class drjit.cuda.ArrayXi#

Bases: ArrayBase

class drjit.cuda.Array0f64#

Bases: ArrayBase

class drjit.cuda.Array1f64#

Bases: ArrayBase

class drjit.cuda.Array2f64#

Bases: ArrayBase

class drjit.cuda.Array3f64#

Bases: ArrayBase

class drjit.cuda.Array4f64#

Bases: ArrayBase

class drjit.cuda.ArrayXf64#

Bases: ArrayBase

class drjit.cuda.Array0u64#

Bases: ArrayBase

class drjit.cuda.Array1u64#

Bases: ArrayBase

class drjit.cuda.Array2u64#

Bases: ArrayBase

class drjit.cuda.Array3u64#

Bases: ArrayBase

class drjit.cuda.Array4u64#

Bases: ArrayBase

class drjit.cuda.ArrayXu64#

Bases: ArrayBase

class drjit.cuda.Array0i64#

Bases: ArrayBase

class drjit.cuda.Array1i64#

Bases: ArrayBase

class drjit.cuda.Array2i64#

Bases: ArrayBase

class drjit.cuda.Array3i64#

Bases: ArrayBase

class drjit.cuda.Array4i64#

Bases: ArrayBase

class drjit.cuda.ArrayXi64#

Bases: ArrayBase

class drjit.cuda.Matrix2u#

Bases: ArrayBase

class drjit.cuda.Matrix2i#

Bases: ArrayBase

class drjit.cuda.Matrix2f#

Bases: ArrayBase

class drjit.cuda.Matrix2f64#

Bases: ArrayBase

class drjit.cuda.Matrix3u#

Bases: ArrayBase

class drjit.cuda.Matrix3i#

Bases: ArrayBase

class drjit.cuda.Matrix3f#

Bases: ArrayBase

class drjit.cuda.Matrix3f64#

Bases: ArrayBase

class drjit.cuda.Matrix4u#

Bases: ArrayBase

class drjit.cuda.Matrix4i#

Bases: ArrayBase

class drjit.cuda.Matrix4f#

Bases: ArrayBase

class drjit.cuda.Matrix4f64#

Bases: ArrayBase

class drjit.cuda.Complex2f#

Bases: ArrayBase

class drjit.cuda.Complex2f64#

Bases: ArrayBase

class drjit.cuda.Quaternion4f#

Bases: ArrayBase

class drjit.cuda.Quaternion4f64#

Bases: ArrayBase

CUDA array namespace with automatic differentiation (drjit.cuda.ad)#

class drjit.cuda.ad.Bool#

Bases: ArrayBase

class drjit.cuda.ad.Float#

Bases: ArrayBase

class drjit.cuda.ad.Float64#

Bases: ArrayBase

class drjit.cuda.ad.UInt#

Bases: ArrayBase

class drjit.cuda.ad.UInt64#

Bases: ArrayBase

class drjit.cuda.ad.Int#

Bases: ArrayBase

class drjit.cuda.ad.Int64#

Bases: ArrayBase

class drjit.cuda.ad.Array0b#

Bases: ArrayBase

class drjit.cuda.ad.Array1b#

Bases: ArrayBase

class drjit.cuda.ad.Array2b#

Bases: ArrayBase

class drjit.cuda.ad.Array3b#

Bases: ArrayBase

class drjit.cuda.ad.Array4b#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXb#

Bases: ArrayBase

class drjit.cuda.ad.Array0f#

Bases: ArrayBase

class drjit.cuda.ad.Array1f#

Bases: ArrayBase

class drjit.cuda.ad.Array2f#

Bases: ArrayBase

class drjit.cuda.ad.Array3f#

Bases: ArrayBase

class drjit.cuda.ad.Array4f#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXf#

Bases: ArrayBase

class drjit.cuda.ad.Array0u#

Bases: ArrayBase

class drjit.cuda.ad.Array1u#

Bases: ArrayBase

class drjit.cuda.ad.Array2u#

Bases: ArrayBase

class drjit.cuda.ad.Array3u#

Bases: ArrayBase

class drjit.cuda.ad.Array4u#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXu#

Bases: ArrayBase

class drjit.cuda.ad.Array0i#

Bases: ArrayBase

class drjit.cuda.ad.Array1i#

Bases: ArrayBase

class drjit.cuda.ad.Array2i#

Bases: ArrayBase

class drjit.cuda.ad.Array3i#

Bases: ArrayBase

class drjit.cuda.ad.Array4i#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXi#

Bases: ArrayBase

class drjit.cuda.ad.Array0f64#

Bases: ArrayBase

class drjit.cuda.ad.Array1f64#

Bases: ArrayBase

class drjit.cuda.ad.Array2f64#

Bases: ArrayBase

class drjit.cuda.ad.Array3f64#

Bases: ArrayBase

class drjit.cuda.ad.Array4f64#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXf64#

Bases: ArrayBase

class drjit.cuda.ad.Array0u64#

Bases: ArrayBase

class drjit.cuda.ad.Array1u64#

Bases: ArrayBase

class drjit.cuda.ad.Array2u64#

Bases: ArrayBase

class drjit.cuda.ad.Array3u64#

Bases: ArrayBase

class drjit.cuda.ad.Array4u64#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXu64#

Bases: ArrayBase

class drjit.cuda.ad.Array0i64#

Bases: ArrayBase

class drjit.cuda.ad.Array1i64#

Bases: ArrayBase

class drjit.cuda.ad.Array2i64#

Bases: ArrayBase

class drjit.cuda.ad.Array3i64#

Bases: ArrayBase

class drjit.cuda.ad.Array4i64#

Bases: ArrayBase

class drjit.cuda.ad.ArrayXi64#

Bases: ArrayBase

class drjit.cuda.ad.Matrix2u#

Bases: ArrayBase

class drjit.cuda.ad.Matrix2i#

Bases: ArrayBase

class drjit.cuda.ad.Matrix2f#

Bases: ArrayBase

class drjit.cuda.ad.Matrix2f64#

Bases: ArrayBase

class drjit.cuda.ad.Matrix3u#

Bases: ArrayBase

class drjit.cuda.ad.Matrix3i#

Bases: ArrayBase

class drjit.cuda.ad.Matrix3f#

Bases: ArrayBase

class drjit.cuda.ad.Matrix3f64#

Bases: ArrayBase

class drjit.cuda.ad.Matrix4u#

Bases: ArrayBase

class drjit.cuda.ad.Matrix4i#

Bases: ArrayBase

class drjit.cuda.ad.Matrix4f#

Bases: ArrayBase

class drjit.cuda.ad.Matrix4f64#

Bases: ArrayBase

class drjit.cuda.ad.Complex2f#

Bases: ArrayBase

class drjit.cuda.ad.Complex2f64#

Bases: ArrayBase

class drjit.cuda.ad.Quaternion4f#

Bases: ArrayBase

class drjit.cuda.ad.Quaternion4f64#

Bases: ArrayBase

LLVM array namespace (drjit.llvm)#

class drjit.llvm.Bool#

Bases: ArrayBase

class drjit.llvm.Float#

Bases: ArrayBase

class drjit.llvm.Float64#

Bases: ArrayBase

class drjit.llvm.UInt#

Bases: ArrayBase

class drjit.llvm.UInt64#

Bases: ArrayBase

class drjit.llvm.Int#

Bases: ArrayBase

class drjit.llvm.Int64#

Bases: ArrayBase

class drjit.llvm.Array0b#

Bases: ArrayBase

class drjit.llvm.Array1b#

Bases: ArrayBase

class drjit.llvm.Array2b#

Bases: ArrayBase

class drjit.llvm.Array3b#

Bases: ArrayBase

class drjit.llvm.Array4b#

Bases: ArrayBase

class drjit.llvm.ArrayXb#

Bases: ArrayBase

class drjit.llvm.Array0f#

Bases: ArrayBase

class drjit.llvm.Array1f#

Bases: ArrayBase

class drjit.llvm.Array2f#

Bases: ArrayBase

class drjit.llvm.Array3f#

Bases: ArrayBase

class drjit.llvm.Array4f#

Bases: ArrayBase

class drjit.llvm.ArrayXf#

Bases: ArrayBase

class drjit.llvm.Array0u#

Bases: ArrayBase

class drjit.llvm.Array1u#

Bases: ArrayBase

class drjit.llvm.Array2u#

Bases: ArrayBase

class drjit.llvm.Array3u#

Bases: ArrayBase

class drjit.llvm.Array4u#

Bases: ArrayBase

class drjit.llvm.ArrayXu#

Bases: ArrayBase

class drjit.llvm.Array0i#

Bases: ArrayBase

class drjit.llvm.Array1i#

Bases: ArrayBase

class drjit.llvm.Array2i#

Bases: ArrayBase

class drjit.llvm.Array3i#

Bases: ArrayBase

class drjit.llvm.Array4i#

Bases: ArrayBase

class drjit.llvm.ArrayXi#

Bases: ArrayBase

class drjit.llvm.Array0f64#

Bases: ArrayBase

class drjit.llvm.Array1f64#

Bases: ArrayBase

class drjit.llvm.Array2f64#

Bases: ArrayBase

class drjit.llvm.Array3f64#

Bases: ArrayBase

class drjit.llvm.Array4f64#

Bases: ArrayBase

class drjit.llvm.ArrayXf64#

Bases: ArrayBase

class drjit.llvm.Array0u64#

Bases: ArrayBase

class drjit.llvm.Array1u64#

Bases: ArrayBase

class drjit.llvm.Array2u64#

Bases: ArrayBase

class drjit.llvm.Array3u64#

Bases: ArrayBase

class drjit.llvm.Array4u64#

Bases: ArrayBase

class drjit.llvm.ArrayXu64#

Bases: ArrayBase

class drjit.llvm.Array0i64#

Bases: ArrayBase

class drjit.llvm.Array1i64#

Bases: ArrayBase

class drjit.llvm.Array2i64#

Bases: ArrayBase

class drjit.llvm.Array3i64#

Bases: ArrayBase

class drjit.llvm.Array4i64#

Bases: ArrayBase

class drjit.llvm.ArrayXi64#

Bases: ArrayBase

class drjit.llvm.Matrix2u#

Bases: ArrayBase

class drjit.llvm.Matrix2i#

Bases: ArrayBase

class drjit.llvm.Matrix2f#

Bases: ArrayBase

class drjit.llvm.Matrix2f64#

Bases: ArrayBase

class drjit.llvm.Matrix3u#

Bases: ArrayBase

class drjit.llvm.Matrix3i#

Bases: ArrayBase

class drjit.llvm.Matrix3f#

Bases: ArrayBase

class drjit.llvm.Matrix3f64#

Bases: ArrayBase

class drjit.llvm.Matrix4u#

Bases: ArrayBase

class drjit.llvm.Matrix4i#

Bases: ArrayBase

class drjit.llvm.Matrix4f#

Bases: ArrayBase

class drjit.llvm.Matrix4f64#

Bases: ArrayBase

class drjit.llvm.Complex2f#

Bases: ArrayBase

class drjit.llvm.Complex2f64#

Bases: ArrayBase

class drjit.llvm.Quaternion4f#

Bases: ArrayBase

class drjit.llvm.Quaternion4f64#

Bases: ArrayBase

LLVM array namespace with automatic differentiation (drjit.llvm.ad)#

class drjit.llvm.ad.Bool#

Bases: ArrayBase

class drjit.llvm.ad.Float#

Bases: ArrayBase

class drjit.llvm.ad.Float64#

Bases: ArrayBase

class drjit.llvm.ad.UInt#

Bases: ArrayBase

class drjit.llvm.ad.UInt64#

Bases: ArrayBase

class drjit.llvm.ad.Int#

Bases: ArrayBase

class drjit.llvm.ad.Int64#

Bases: ArrayBase

class drjit.llvm.ad.Array0b#

Bases: ArrayBase

class drjit.llvm.ad.Array1b#

Bases: ArrayBase

class drjit.llvm.ad.Array2b#

Bases: ArrayBase

class drjit.llvm.ad.Array3b#

Bases: ArrayBase

class drjit.llvm.ad.Array4b#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXb#

Bases: ArrayBase

class drjit.llvm.ad.Array0f#

Bases: ArrayBase

class drjit.llvm.ad.Array1f#

Bases: ArrayBase

class drjit.llvm.ad.Array2f#

Bases: ArrayBase

class drjit.llvm.ad.Array3f#

Bases: ArrayBase

class drjit.llvm.ad.Array4f#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXf#

Bases: ArrayBase

class drjit.llvm.ad.Array0u#

Bases: ArrayBase

class drjit.llvm.ad.Array1u#

Bases: ArrayBase

class drjit.llvm.ad.Array2u#

Bases: ArrayBase

class drjit.llvm.ad.Array3u#

Bases: ArrayBase

class drjit.llvm.ad.Array4u#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXu#

Bases: ArrayBase

class drjit.llvm.ad.Array0i#

Bases: ArrayBase

class drjit.llvm.ad.Array1i#

Bases: ArrayBase

class drjit.llvm.ad.Array2i#

Bases: ArrayBase

class drjit.llvm.ad.Array3i#

Bases: ArrayBase

class drjit.llvm.ad.Array4i#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXi#

Bases: ArrayBase

class drjit.llvm.ad.Array0f64#

Bases: ArrayBase

class drjit.llvm.ad.Array1f64#

Bases: ArrayBase

class drjit.llvm.ad.Array2f64#

Bases: ArrayBase

class drjit.llvm.ad.Array3f64#

Bases: ArrayBase

class drjit.llvm.ad.Array4f64#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXf64#

Bases: ArrayBase

class drjit.llvm.ad.Array0u64#

Bases: ArrayBase

class drjit.llvm.ad.Array1u64#

Bases: ArrayBase

class drjit.llvm.ad.Array2u64#

Bases: ArrayBase

class drjit.llvm.ad.Array3u64#

Bases: ArrayBase

class drjit.llvm.ad.Array4u64#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXu64#

Bases: ArrayBase

class drjit.llvm.ad.Array0i64#

Bases: ArrayBase

class drjit.llvm.ad.Array1i64#

Bases: ArrayBase

class drjit.llvm.ad.Array2i64#

Bases: ArrayBase

class drjit.llvm.ad.Array3i64#

Bases: ArrayBase

class drjit.llvm.ad.Array4i64#

Bases: ArrayBase

class drjit.llvm.ad.ArrayXi64#

Bases: ArrayBase

class drjit.llvm.ad.Matrix2u#

Bases: ArrayBase

class drjit.llvm.ad.Matrix2i#

Bases: ArrayBase

class drjit.llvm.ad.Matrix2f#

Bases: ArrayBase

class drjit.llvm.ad.Matrix2f64#

Bases: ArrayBase

class drjit.llvm.ad.Matrix3u#

Bases: ArrayBase

class drjit.llvm.ad.Matrix3i#

Bases: ArrayBase

class drjit.llvm.ad.Matrix3f#

Bases: ArrayBase

class drjit.llvm.ad.Matrix3f64#

Bases: ArrayBase

class drjit.llvm.ad.Matrix4u#

Bases: ArrayBase

class drjit.llvm.ad.Matrix4i#

Bases: ArrayBase

class drjit.llvm.ad.Matrix4f#

Bases: ArrayBase

class drjit.llvm.ad.Matrix4f64#

Bases: ArrayBase

class drjit.llvm.ad.Complex2f#

Bases: ArrayBase

class drjit.llvm.ad.Complex2f64#

Bases: ArrayBase

class drjit.llvm.ad.Quaternion4f#

Bases: ArrayBase

class drjit.llvm.ad.Quaternion4f64#

Bases: ArrayBase