def fmadd_(a0, a1, a2):
if not a0.IsFloat:
raise Exception("fmadd(): requires floating point operands!")
if not a0.IsSpecial:
ar, sr = _check3(a0, a1, a2)
for i in range(sr):
ar[i] = _ek.fmadd(a0[i], a1[i], a2[i])
return ar
else:
return a0 * a1 + a2
def imul_(a0, a1):
if not a0.IsArithmetic:
raise Exception("imul(): requires arithmetic operands!")
sr = _check2_inplace(a0, a1)
if not a0.IsSpecial:
for i in range(sr):
a0[i] *= a1[i]
else:
a0.assign(a0 * a1)
return a0
def rsqrt_(a0):
if not a0.IsFloat:
raise Exception("rsqrt(): requires floating point operands!")
if not a0.IsSpecial:
ar, sr = _check1(a0)
for i in range(sr):
ar[i] = _ek.rsqrt(a0[i])
return ar
else:
return _ek.rcp(_ek.sqrt(a0))
def not_(a0):
if a0.IsFloat:
raise Exception("not(): requires an integral or mask operand!")
ar, sr = _check1(a0)
if a0.Value is bool:
for i in range(sr):
ar[i] = not a0[i]
else:
for i in range(sr):
ar[i] = ~a0[i]
return ar
def log_(a0):
if not a0.IsFloat:
raise Exception("log(): requires floating point operands!")
ar, sr = _check1(a0)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = _ek.log(a0[i])
elif a0.IsComplex:
ar.real = .5 * _ek.log(_ek.squared_norm(a0))
ar.imag = _ek.arg(a0)
elif a0.IsQuaternion:
qi_n = _ek.normalize(a0.imag)
rq = _ek.norm(a0)
acos_rq = _ek.acos(a0.real / rq)
log_rq = _ek.log(rq)
ar.imag = qi_n * acos_rq
ar.real = log_rq
else:
raise Exception("log(): unsupported array type!")
return ar
def grad_suspended_(a):
if not a.IsDiff:
raise Exception("Expected a differentiable array type!")
if a.Depth > 1:
suspended = False
for i in range(len(a)):
suspended |= a[i].grad_suspended_()
return suspended
else:
return a.index() < 0
def sincosh_(a0):
if not a0.IsFloat:
raise Exception("sincosh(): requires floating point operands!")
ar0, sr0 = _check1(a0)
ar1 = a0.empty_(sr0 if a0.Size == Dynamic else 0)
if not a0.IsSpecial:
for i in range(sr0):
result = _ek.sincosh(a0[i])
ar0[i] = result[0]
ar1[i] = result[1]
elif a0.IsComplex:
s, c = _ek.sincos(a0.imag)
sh, ch = _ek.sincosh(a0.real)
ar0.real = sh * c
ar0.imag = ch * s
ar1.real = ch * c
ar1.imag = sh * s
else:
raise Exception("sincosh(): unsupported array type!")
return ar0, ar1
def numpy(a):
import numpy
arr = numpy.array(a, copy=False)
if a.IsComplex:
arr = numpy.ascontiguousarray(arr)
if arr.dtype == numpy.float32:
return arr.view(numpy.complex64)[..., 0]
elif arr.dtype == numpy.float64:
return arr.view(numpy.complex128)[..., 0]
else:
raise Exception("Unsupported dtype for complex conversion!")
return arr
def itruediv_(a0, a1):
if not a0.IsFloat:
raise Exception("Use the floor division operator \"//=\" for "
"Enoki integer arrays.")
sr = _check2_inplace(a0, a1)
if not a0.IsSpecial:
for i in range(sr):
a0[i] /= a1[i]
else:
a0.assign(a0 / a1)
return a0
def exp_(a0):
if not a0.IsFloat:
raise Exception("exp(): requires floating point operands!")
ar, sr = _check1(a0)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = _ek.exp(a0[i])
elif a0.IsComplex:
s, c = _ek.sincos(a0.imag)
exp_r = _ek.exp(a0.real)
ar.real = exp_r * c
ar.imag = exp_r * s
elif a0.IsQuaternion:
qi = a0.imag
ri = _ek.norm(qi)
exp_w = _ek.exp(a0.real)
s, c = _ek.sincos(ri)
ar.imag = qi * (s * exp_w / ri)
ar.real = c * exp_w
else:
raise Exception("exp(): unsupported array type!")
return ar
def diff_array_t(a):
if isinstance(a, tuple):
return tuple(diff_array_t(v) for v in a)
elif isinstance(a, list):
return [diff_array_t(v) for v in a]
elif not is_array_v(a):
raise Exception("diff_array_t(): requires an Enoki input array!")
elif not isinstance(a, type):
return diff_array_t(type(a))(a)
elif a.IsDiff:
return a
else:
return a.ReplaceScalar(a.Type, diff=True)
def mul_(a0, a1):
if not a0.IsArithmetic:
raise Exception("mul(): requires arithmetic operands!")
if _ek.array_depth_v(a1) < _ek.array_depth_v(a0):
# Avoid type promotion in scalars multiplication, which would
# be costly for special types (matrices, quaternions, etc.)
sr = len(a0)
ar = a0.empty_(sr if a0.Size == Dynamic else 0)
for i in range(len(a0)):
ar[i] = a0[i] * a1
return ar
ar, sr = _check2(a0, a1)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = a0[i] * a1[i]
elif a0.IsComplex:
ar.real = _ek.fmsub(a0.real, a1.real, a0.imag * a1.imag)
ar.imag = _ek.fmadd(a0.real, a1.imag, a0.imag * a1.real)
elif a0.IsQuaternion:
tbl = (4, 3, -2, 1, -3, 4, 1, 2, 2, -1, 4, 3, -1, -2, -3, 4)
for i in range(4):
accum = 0
for j in range(4):
idx = tbl[i * 4 + j]
value = a1[abs(idx) - 1]
accum = _ek.fmadd(a0[j], value if idx > 0 else -value, accum)
ar[i] = accum
elif a0.IsMatrix:
raise Exception(
"mul(): please use the matrix multiplication operator '@' instead."
)
else:
raise Exception("mul(): unsupported array type!")
return ar
def sqrt_(a0):
if not a0.IsFloat:
raise Exception("sqrt(): requires floating point operands!")
ar, sr = _check1(a0)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = _ek.sqrt(a0[i])
elif a0.IsComplex:
n = abs(a0)
m = a0.real >= 0
zero = _ek.eq(n, 0)
t1 = _ek.sqrt(.5 * (n + abs(a0.real)))
t2 = .5 * a0.imag / t1
im = _ek.select(m, t2, _ek.copysign(t1, a0.imag))
ar.real = _ek.select(m, t1, abs(t2))
ar.imag = _ek.select(zero, 0, im)
elif a0.IsQuaternion:
ri = _ek.norm(a0.imag)
cs = _ek.sqrt(a0.Complex(a0.real, ri))
ar.imag = a0.imag * (_ek.rcp(ri) * cs.imag)
ar.real = cs.real
else:
raise Exception("sqrt(): unsupported array type!")
return ar
def pow_(a0, a1):
if not a0.IsFloat:
raise Exception("pow(): requires floating point operands!")
if not a0.IsSpecial:
if isinstance(a1, int) or isinstance(a1, float):
ar, sr = _check1(a0)
for i in range(sr):
ar[i] = _ek.pow(a0[i], a1)
else:
ar, sr = _check2(a0, a1)
for i in range(sr):
ar[i] = _ek.pow(a0[i], a1[i])
else:
return _ek.exp(_ek.log(a0) * a1)
return ar
def float_array_t(a):
if not is_array_v(a):
return int
size = a.Type.Size
if size == 2:
vt = VarType.Float16
elif size == 4:
vt = VarType.Float32
elif size == 8:
vt = VarType.Float64
else:
raise Exception("Unsupported variable size!")
t = a.ReplaceScalar(vt)
return t if isinstance(a, type) else t(a)
def broadcast_(self, value):
if not self.IsSpecial:
for i in range(len(self)):
self.set_entry_(i, value)
elif self.IsComplex:
self.set_entry_(0, value)
self.set_entry_(1, 0)
elif self.IsQuaternion:
for i in range(3):
self.set_entry_(i, 0)
self.set_entry_(3, value)
elif self.IsMatrix:
t = self.Value
for i in range(len(self)):
c = _ek.zero(t)
c.set_entry_(i, t.Value(value))
self.set_entry_(i, c)
else:
raise Exception("broadcast_(): don't know how to handle this type!")
def index_ad_(a):
if not a.IsDiff:
raise Exception("Expected a differentiable array type!")
return tuple(v.index_ad() for v in a)
def set_grad_suspended_(a, value):
if not a.IsDiff:
raise Exception("Expected a differentiable array type!")
for i in range(len(a)):
a.entry_ref_(i).set_grad_suspended_(value)
return a
def index_(a):
if not a.IsJIT:
raise Exception("Expected a JIT array type!")
return tuple(v.index() for v in a)
def enqueue_(a):
if not a.IsDiff:
raise Exception("Expected a differentiable array type!")
for i in range(len(a)):
a[i].enqueue_()
def op_cuda_array_interface(a):
if not a.IsCUDA:
raise Exception("__cuda_array_interface__: only CUDA "
"arrays are supported!")
return a.export_(migrate_to_host=False, version=2)
