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
. Whenshape
specifies a sequence, it must be compatible with static dimensions of thedtype
. For example,dr.zeros(dr.cuda.Array2f, shape=(3, 100))
fails, since the leading dimension is incompatible withdrjit.cuda.Array2f
. Whenshape
is an integer, it specifies the size of the last (dynamic) dimension, if available.A tensorial type like
drjit.scalar.TensorXf
. Whenshape
specifies a sequence (list/tuple/..), it determines the tensor rank and shape. Whenshape
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
, orbool
. Theshape
parameter is ignored in this case.
Note that when
dtype
refers to a scalar mask or a mask array, it will be initialized toFalse
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:A Dr.Jit array type like
drjit.cuda.Array2f
. Whenshape
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 withdrjit.cuda.Array2f
. Whenshape
is an integer, it specifiesthe size of the last (dynamic) dimension, if available.
A tensorial type like
drjit.scalar.TensorXf
. Whenshape
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
, orbool
. Theshape
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:A Dr.Jit array type like
drjit.cuda.Array2f
. Whenshape
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 withdrjit.cuda.Array2f
. Whenshape
is an integer, it specifies the size of the last (dynamic) dimension, if available.A tensorial type like
drjit.scalar.TensorXf
. Whenshape
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
, orbool
. Theshape
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 withvalue
- 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
, orbool
. Thesize
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 sizesize
- 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 sizestep
, wherestart
= 0 andstep
= 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 asdrjit.scalar.ArrayXu
ordrjit.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 asdrjit.scalar.ArrayXf
ordrjit.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 positionindex
. It expects adtype
argument and will return an instance of this type. The optionalactive
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 likedrjit.llvm.ad.Float
, this operation implements a parallelized version of the Python array indexing expressionsource[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)
- 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
- When
type(source)
matchesdtype
, 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) )
- When
- Otherwise, the operation reconstructs the requested
dtype
from a flat source
array, using C-style ordering with a suitably modifiedindex
. 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) )
- Otherwise, the operation reconstructs the requested
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
ordrjit.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
ordrjit.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 thetarget
array at positionindex
. The optionalactive
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 likedrjit.llvm.ad.Float
, this operation implements a parallelized version of the Python array indexing expressiontarget[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
andvalue
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)
- When
- 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)
- Otherwise, the operation flattens the
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
ordrjit.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
ordrjit.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 thetarget
array at positionindex
. The optionalactive
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 (seedrjit.ReduceOp
).This operation can be used in the following different ways:
1. When
target
is a 1D Dr.Jit array likedrjit.llvm.ad.Float
, this operation implements a parallelized version of the expressiontarget[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.), thentarget
andvalue
must be of the same type. The scatter reduce operation threads through entries and invokes itself recursively. For example, the scatter operation intarget = 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
ordrjit.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
ordrjit.cuda.Bool
) specifying active components. Dr.Jit will attempt an implicit conversion if another type is provided. The default is True.
- 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]
withorder='F'
(the default for Dr.Jit arrays),
which means that X/Y/Z components alternate. -
[1, 2, 3, 4, 5, 6]
withorder='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. Whenarray
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 todrjit.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 withorder='F'
(the default for Dr.Jit arrays), and[1, 3, 5]
and[2, 4, 6]
when unraveled withorder='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.
- An instance of type
- 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 Pythonlist
. For custom source structure the returned type needs to be specified through the argumentreturn_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 of1
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:
arg0 (list | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type
arg1 (list | drjit.ArrayBase) – A Python or Dr.Jit arithmetic type
- 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:
perm (drjit.ArrayBase) – A Python list of integers
value (drjit.ArrayBase) – A Dr.Jit static array type
- Returns:
Shuffled input array
- drjit.compress(arg, /) int | drjit.ArrayBase #
Compress a mask into a array of nonzero indices.
This function takes an boolean array as input and then returns the indices of nonzero entries.
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 typeint
. 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 arrayvalue
and writes them to the returned array. For example,a, b, c, d, e, f
turns intoa+b, c+d, e+f
whensize == 2
. The length of the input array must be a multiple ofsize
.- 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:
condition (bool | drjit.ArrayBase) – A Python or Dr.Jit mask/boolean array
x (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
y (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
- 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
andb
- 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:
- 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:
- 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:
- 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)
andarg.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
- drjit.binary_search(start, end, pred)#
Perform binary search over a range given a predicate
pred
, which monotonically decreases over this range (i.e. max oneTrue
->False
transition).Given a (scalar)
start
andend
index of a range, this function evaluates a predicatefloor(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 parameterIndex
is automatically inferred from the supplied predicate. Specifically, the predicate takes an index array as input argument. Whenpred
isFalse
for all entries, the function returnsstart
, and when it isTrue
for all cases, it returnsend
.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 ofx
to the first indexj
such thatdata[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)#
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 targetshape
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. Targetshape
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)
- 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
andb
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 whatdrjit.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 todrjit.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 todrjit.schedule()
. Variables that are garbage collected between the originaldrjit.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 callsdrjit.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
ifarg
or type(arg
) is a Dr.Jit array type, andFalse
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
ifarg
represents a Dr.Jit mask array or Pythonbool
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
ifarg
represents a Dr.Jit floating point array or Pythonfloat
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
ifarg
represents an integral Dr.Jit array or Pythonint
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
orfloat
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
ifarg
represents an arithmetic Dr.Jit array or Pythonint
orfloat
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
orfloat
value/type.- Parameters:
arg (object) – An arbitrary Python object
- Returns:
True
ifarg
represents an signed Dr.Jit array or Pythonint
orfloat
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
ifarg
represents an unsigned Dr.Jit array or Pythonbool
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
ifarg
represents an array type from thedrjit.cuda.*
ordrjit.llvm.*
namespaces, andFalse
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
ifarg
represents an array type from thedrjit.llvm.*
namespace, andFalse
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 satisfiesis_diff_v(..) = 1
.- Parameters:
arg (object) – An arbitrary Python object
- Returns:
True
ifarg
represents an array type from thedrjit.[cuda/llvm].ad.*
namespace, andFalse
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, andFalse
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, andFalse
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, andFalse
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, andFalse
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, andFalse
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
ifarg
has aDRJIT_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
whenarg
is a dynamic Dr.Jit array. Returns1
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 depth1
, while an array consisting of sub-arrays (e.g.,drjit.cuda.Array3f
) has depth2
.- 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:
value (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
min (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
max (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
- 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:
value (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
min (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
max (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
- Returns:
Clipped input
- Return type:
float | drjit.ArrayBase
- drjit.fma(arg0, arg1, arg2, /)#
Perform a fused multiply-add (FMA) operation.
Given arguments
arg0
,arg1
, andarg2
, this operation computesarg0
*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:
arg0 (float | drjit.ArrayBase) – First multiplication operand
arg1 (float | drjit.ArrayBase) – Second multiplication operand
arg2 (float | drjit.ArrayBase) – Additive operand
- 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
andb
using the expression \(a \cdot (1 - t) + b \cdot t\)- Parameters:
a (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
b (int | float | drjit.ArrayBase) – A Python or Dr.Jit array
t (float | drjit.ArrayBase) – A Python or Dr.Jit array
- 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 ofarg1
- 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 ofarg1
- 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
ordrjit.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
ordrjit.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
ordrjit.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:
arg0 (list | drjit.ArrayBase) – A Python or Dr.Jit 3D type
arg1 (list | drjit.ArrayBase) – A Python or Dr.Jit 3D type
- 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:
y (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type
x (float | drjit.ArrayBase) – A Python or Dr.Jit floating point type
- 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 whenself
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 ofsrc
.Broadcasting is applied to
dst
if necessary to match the type ofsrc
.- 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 nothingADFlag.ClearEdges
: delete all traversed edges from the computation graphADFlag.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 intactADFlag.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:
Discover and enqueue the variables to be considered as inputs during the subsequent AD traversal.
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 providedvalue
will be considered as input during the subsequent AD traversal.ADMode.Backward
: a traversal of the AD graph starting from the providedvalue
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
tob
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 useADMode.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 thesource
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#