def slice_update_train_job(
x: otp.Numpy.Placeholder(shape=input_shape, dtype=dtype),
update: otp.Numpy.Placeholder(shape=update_shape, dtype=dtype),
) -> otp.Numpy:
x_var = flow.get_variable(
shape=input_shape,
dtype=dtype,
initializer=flow.constant_initializer(0.0),
name="x",
)
update_var = flow.get_variable(
shape=update_shape,
dtype=dtype,
initializer=flow.constant_initializer(0.0),
name="update",
)
x = x + x_var
update = update + update_var
if callable(diff_watcher_maker):
flow.watch_diff(x, diff_watcher_maker(input_shape))
flow.watch_diff(update, diff_watcher_maker(update_shape))
y = flow.slice_update(x, update, slice_tup_list)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[1e-3]),
momentum=0).minimize(y)
return y
def test_slice_update(test_case):
x = np.array([1, 1, 1, 1, 1]).astype(np.float32)
input = flow.tensor(x, requires_grad=True)
update = flow.tensor(np.array([2, 3, 4]).astype(np.float32),
requires_grad=True)
output = np.array([1.0, 2.0, 3.0, 4.0, 1.0])
y = flow.slice_update(input, update, slice_tup_list=[[1, 4, 1]])
z = y.sum()
z.backward()
test_case.assertTrue(np.array_equal(y.numpy(), output))
np_grad = np.zeros(x.shape)
np_grad[0] = 1
np_grad[4] = 1
test_case.assertTrue(np.array_equal(input.grad.numpy(), np_grad))
test_case.assertTrue(
np.array_equal(update.grad.numpy(), np.ones(update.shape)))
def slice_update_job(
x: otp.Numpy.Placeholder(shape=input_shape, dtype=dtype),
update: otp.Numpy.Placeholder(shape=update_shape, dtype=dtype),
) -> otp.Numpy:
return flow.slice_update(x, update, slice_tup_list)