def where( condition: remote_blob_util.BlobDef, x: Optional[remote_blob_util.BlobDef] = None, y: Optional[remote_blob_util.BlobDef] = None, name: Optional[str] = None, ) -> remote_blob_util.BlobDef: if x is None and y is None: return argwhere(condition, name=name) elif x is not None and y is not None: if name is None: name = id_util.UniqueStr("Where_") if x.shape == condition.shape and y.shape == condition.shape: broadcast_cond = condition broadcast_x = x broadcast_y = y else: broadcast_cond = flow.broadcast_to_compatible_with(condition, [x, y]) broadcast_x = flow.broadcast_to_compatible_with(x, [condition, y]) broadcast_y = flow.broadcast_to_compatible_with(y, [condition, x]) return ( flow.user_op_builder(name) .Op("where") .Input("condition", [broadcast_cond]) .Input("x", [broadcast_x]) .Input("y", [broadcast_y]) .Output("out") .Build() .InferAndTryRun() .RemoteBlobList()[0] ) else: raise ValueError("it is not supported when exactly one of x or y is non-None")
def broadcast_to_compatible_with_fn(x_def: oft.Numpy.Placeholder( x.shape, dtype=flow.float)): x_var = flow.get_variable( "x_var", shape=x.shape, dtype=flow.float, initializer=flow.constant_initializer(0), trainable=True, ) compatible_var = [ flow.get_variable( "compatible_var_{}".format(i), shape=cp_shape, dtype=flow.float, initializer=flow.random_normal_initializer(), trainable=False, ) for i, cp_shape in enumerate(compatible_shape) ] x_var = x_var + x_def y = flow.broadcast_to_compatible_with(x_var, compatible_var) flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0).minimize(y) flow.watch_diff(x_var, dx_watcher) return y
def broadcast_to_compatible_with_fn( x_def: oft.ListNumpy.Placeholder(x_shape, dtype=flow.float), a_def: oft.ListNumpy.Placeholder(a_shape, dtype=flow.float), b_def: oft.ListNumpy.Placeholder(b_shape, dtype=flow.float), ): return flow.broadcast_to_compatible_with( x_def, [flow.identity(a_def), flow.identity(b_def)])
def broadcast_to_compatible_with_fn(x_def: oft.ListNumpy.Placeholder( shape=x_shape, dtype=flow.float)): compatible_var = [ flow.get_variable( "compatible_var_{}".format(i), shape=cp_shape, dtype=flow.float, initializer=flow.random_normal_initializer(), trainable=False, ) for i, cp_shape in enumerate(compatible_shape) ] return flow.broadcast_to_compatible_with(x_def, compatible_var)
def broadcast_to_compatible_with_fn(x_def: oft.Numpy.Placeholder( x.shape, dtype=flow.float)): x_var = flow.get_variable( "x_var", shape=x.shape, dtype=flow.float, initializer=flow.constant_initializer(0), trainable=True, ) compatible_var = [ flow.get_variable( "compatible_var_{}".format(i), shape=cp_shape, dtype=flow.float, initializer=flow.random_normal_initializer(), trainable=False, ) for i, cp_shape in enumerate(compatible_shape) ] x_var = x_var + x_def y = flow.broadcast_to_compatible_with(x_var, compatible_var) flow.losses.add_loss(y) flow.watch_diff(x_var, dx_watcher) return y