def test22_scatter_fwd(m):
x = m.Float(4.0)
ek.enable_grad(x)
values = x * x * ek.linspace(m.Float, 1, 4, 4)
idx = 2 * ek.arange(m.UInt32, 4)
buf = ek.zero(m.Float, 10)
ek.scatter(buf, values, idx)
assert ek.grad_enabled(buf)
ref = [16.0, 0.0, 32.0, 0.0, 48.0, 0.0, 64.0, 0.0, 0.0, 0.0]
assert ek.allclose(buf, ref)
ek.forward(x, retain_graph=True)
grad = ek.grad(buf)
ref_grad = [8.0, 0.0, 16.0, 0.0, 24.0, 0.0, 32.0, 0.0, 0.0, 0.0]
assert ek.allclose(grad, ref_grad)
# Overwrite first value with non-diff value, resulting gradient entry should be 0
y = m.Float(3)
idx = m.UInt32(0)
ek.scatter(buf, y, idx)
ref = [3.0, 0.0, 32.0, 0.0, 48.0, 0.0, 64.0, 0.0, 0.0, 0.0]
assert ek.allclose(buf, ref)
ek.forward(x)
grad = ek.grad(buf)
ref_grad = [0.0, 0.0, 16.0, 0.0, 24.0, 0.0, 32.0, 0.0, 0.0, 0.0]
assert ek.allclose(grad, ref_grad)
def write_indices(indices, *args):
for a in args:
if enoki.is_array_v(a):
if a.Depth > 1:
for i in range(len(a)):
write_indices(indices, a.entry_ref_(i))
elif a.IsDiff:
if enoki.grad_enabled(a) and enoki.flag(
enoki.JitFlag.LoopRecord):
raise enoki.Exception(
"write_indices(): one of the supplied loop "
"variables is attached to the AD graph (i.e. "
"grad_enabled(..) is true). However, recorded "
"loops cannot be differentiated in their entirety. "
"You have two options: either disable loop "
"recording via set_flag(JitFlag.LoopRecord, "
"False). Alternatively, you could implement the "
"adjoint of the loop using ek::CustomOp.")
idx = indices.pop(0)
a.set_index_(idx[0])
a.set_index_ad_(idx[1])
elif a.IsJIT:
idx = indices.pop(0)
a.set_index_(idx[0])
assert idx[1] == 0
elif isinstance(a, tuple) or isinstance(a, list):
for b in a:
if getattr(b, '__name__', None) == '<lambda>':
write_indices(indices, b())
else:
write_indices(indices, b)
elif enoki.is_enoki_struct_v(a):
for k, v in type(a).ENOKI_STRUCT.items():
write_indices(indices, getattr(a, k))
def read_indices(*args):
result = []
for a in args:
if enoki.is_array_v(a):
if a.Depth > 1:
for i in range(len(a)):
result.extend(read_indices(a.entry_ref_(i)))
elif a.IsDiff:
if enoki.grad_enabled(a) and enoki.flag(
enoki.JitFlag.LoopRecord):
raise enoki.Exception(
"read_indices(): one of the supplied loop "
"variables is attached to the AD graph (i.e. "
"grad_enabled(..) is true). However, recorded "
"loops cannot be differentiated in their entirety. "
"You have two options: either disable loop "
"recording via set_flag(JitFlag.LoopRecord, "
"False). Alternatively, you could implement the "
"adjoint of the loop using ek::CustomOp.")
result.append((a.index(), a.index_ad()))
elif a.IsJIT:
result.append((a.index(), 0))
elif isinstance(a, tuple) or isinstance(a, list):
for b in a:
if getattr(b, '__name__', None) == '<lambda>':
result.extend(read_indices(b()))
else:
result.extend(read_indices(b))
elif enoki.is_enoki_struct_v(a):
for k, v in type(a).ENOKI_STRUCT.items():
result.extend(read_indices(getattr(a, k)))
return result
def test_ad_operations(package):
Float, Array3f = package.Float, package.Array3f
prepare(package)
class MyStruct:
ENOKI_STRUCT = {'a': Array3f, 'b': Float}
def __init__(self):
self.a = Array3f()
self.b = Float()
foo = ek.zero(MyStruct, 4)
assert not ek.grad_enabled(foo.a)
assert not ek.grad_enabled(foo.b)
assert not ek.grad_enabled(foo)
ek.enable_grad(foo)
assert ek.grad_enabled(foo.a)
assert ek.grad_enabled(foo.b)
assert ek.grad_enabled(foo)
foo_detached = ek.detach(foo)
assert not ek.grad_enabled(foo_detached.a)
assert not ek.grad_enabled(foo_detached.b)
assert not ek.grad_enabled(foo_detached)
x = Float(4.0)
ek.enable_grad(x)
foo.a += x
foo.b += x * x
ek.forward(x)
foo_grad = ek.grad(foo)
assert foo_grad.a == 1
assert foo_grad.b == 8
ek.set_grad(foo, 5.0)
foo_grad = ek.grad(foo)
assert foo_grad.a == 5.0
assert foo_grad.b == 5.0
ek.accum_grad(foo, 5.0)
foo_grad = ek.grad(foo)
assert foo_grad.a == 10.0
assert foo_grad.b == 10.0
def write_indices(indices, *args):
for a in args:
if enoki.is_array_v(a):
if a.Depth > 1:
for i in range(len(a)):
write_indices(indices, a.entry_ref_(i))
elif a.IsDiff:
if enoki.grad_enabled(a):
raise enoki.Exception(
'Symbolic loop encountered a differentiable array '
'with enabled gradients! This is not supported.')
write_indices(indices, a.detach_())
elif a.IsJIT:
a.set_index_(indices.pop(0))
elif isinstance(a, tuple) or isinstance(a, list):
for b in a:
write_indices(indices, b)
elif enoki.is_enoki_struct_v(a):
for k, v in type(a).ENOKI_STRUCT.items():
write_indices(indices, getattr(a, k))
else:
print(" do not know what to do with %s\n" % str(a))
def test42_suspend_resume(m):
x = m.Array3f(1, 2, 3)
y = m.Array3f(3, 2, 1)
ek.enable_grad(x, y)
assert ek.grad_enabled(x) and ek.grad_enabled(y)
assert not ek.grad_suspended(x) and not ek.grad_suspended(y)
ek.suspend_grad(x, y)
assert not ek.grad_enabled(x) and not ek.grad_enabled(y)
assert ek.grad_suspended(x) and ek.grad_suspended(y)
b = x * y
ek.resume_grad(x, y)
assert ek.grad_enabled(x) and ek.grad_enabled(y)
assert not ek.grad_suspended(x) and not ek.grad_suspended(y)
c = x * y
ek.backward(c)
assert ek.grad(x) == ek.detach(y)
assert ek.grad(y) == ek.detach(x)
ek.suspend_grad(x, y) # validate reference counting of suspended variables
def read_indices(*args):
result = []
for a in args:
if enoki.is_array_v(a):
if a.Depth > 1:
for i in range(len(a)):
result.extend(read_indices(a.entry_ref_(i)))
elif a.IsDiff:
if enoki.grad_enabled(a):
raise enoki.Exception(
'Symbolic loop encountered a differentiable array '
'with enabled gradients! This is not supported.')
result.extend(read_indices(a.detach_()))
elif a.IsJIT:
result.append(a.index())
elif isinstance(a, tuple) or isinstance(a, list):
for b in a:
result.extend(read_indices(b))
elif enoki.is_enoki_struct_v(a):
for k, v in type(a).ENOKI_STRUCT.items():
result.extend(read_indices(getattr(a, k)))
else:
print(" do not know what to do with %s\n" % str(a))
return result