def test_masked_fill_fw_bw_job(
x: oft.Numpy.Placeholder(x_shape, dtype=flow_type),
mask: oft.Numpy.Placeholder(mask_shape, dtype=flow_type),
):
with flow.scope.placement(device, "0:0"):
y = flow.get_variable(
name="vx",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x += flow.cast(y, flow_type)
mask = flow.cast(mask, dtype=flow.int8)
if type_name == "float16":
out = flow.cast(
flow.masked_fill(flow.cast(x, flow.float16), mask, value),
flow.float,
)
else:
out = flow.masked_fill(x, mask, value)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out
def test_fused_scale_tril_fw_bw_job(
x: oft.Numpy.Placeholder(shape, dtype=flow_type), ):
with flow.scope.placement(device, "0:0"):
x_var = flow.get_variable(
name="xv",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x += flow.cast(x_var, dtype=flow_type)
if type_name == "float16":
out = flow.cast(
flow.math.fused_scale_tril(flow.cast(x, flow.float16),
diagonal,
scale=scale),
flow.float,
)
else:
out = flow.math.fused_scale_tril(x, diagonal, scale=scale)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out
def test_element_wise_mul_job(
x: oft.Numpy.Placeholder(shape, dtype=flow.float),
y: oft.Numpy.Placeholder(shape, dtype=flow.float),
):
with flow.scope.placement(device, "0:0"):
x += flow.get_variable(
name="vx",
shape=(1,),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
y += flow.get_variable(
name="vy",
shape=(1,),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x = flow.cast(x, dtype=flow_type)
y = flow.cast(y, dtype=flow_type)
out = flow.math.multiply(x, y)
out = flow.cast(out, dtype=flow.float)
flow.losses.add_loss(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(y, test_global_storage.Setter("y"))
flow.watch_diff(y, test_global_storage.Setter("y_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out
def test_element_wise_mul_job(
x: oft.Numpy.Placeholder(shape, dtype=flow.float),
y: oft.Numpy.Placeholder(shape, dtype=flow.float),
):
with flow.scope.placement(device, "0:0"):
x += flow.get_variable(
name="vx",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
y += flow.get_variable(
name="vy",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x = flow.cast(x, dtype=flow_type)
y = flow.cast(y, dtype=flow_type)
out = flow.math.multiply(x, y)
out = flow.cast(out, dtype=flow.float)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(y, test_global_storage.Setter("y"))
flow.watch_diff(y, test_global_storage.Setter("y_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out