def accumulate_client_votes(vote_accumulator, example):
"""Accumulates client votes on prefix extensions."""
example = tf.strings.lower(example)
# Append the default terminator to the example.
default_terminator = tf.constant(DEFAULT_TERMINATOR, dtype=tf.string)
example = tf.strings.join([example, default_terminator])
# Compute effective round number.
effective_round_num = tf.math.floordiv(round_num, num_sub_rounds)
if tf.strings.length(example) < effective_round_num:
return vote_accumulator
else:
discovered_prefixes_index = discovered_prefixes_table.lookup(
tf.strings.substr(example, 0, effective_round_num))
possible_prefix_extensions_index = possible_prefix_extensions_table.lookup(
tf.strings.substr(example, effective_round_num, 1))
if tf.math.equal(possible_prefix_extensions_index,
tf.constant(DEFAULT_VALUE)):
return vote_accumulator
elif tf.math.equal(discovered_prefixes_index, tf.constant(DEFAULT_VALUE)):
indices = [[0, possible_prefix_extensions_index]]
updates = tf.constant([1])
return tf.tensor_scatter_nd_add(vote_accumulator, indices, updates)
else:
indices = [[
discovered_prefixes_index, possible_prefix_extensions_index
]]
updates = tf.constant([1])
return tf.tensor_scatter_nd_add(vote_accumulator, indices, updates)
def jet_reco(px, py, jet_idx, max_jets):
tf.debugging.assert_shapes(
[
(px, ("N")),
(py, ("N")),
(jet_idx, ("N")),
]
)
jet_idx_capped = tf.where(jet_idx <= max_jets, jet_idx, 0)
jet_px = tf.zeros(
[
max_jets,
],
dtype=px.dtype,
)
jet_py = tf.zeros(
[
max_jets,
],
dtype=py.dtype,
)
jet_px_new = tf.tensor_scatter_nd_add(jet_px, indices=tf.expand_dims(jet_idx_capped, axis=-1), updates=px)
jet_py_new = tf.tensor_scatter_nd_add(jet_py, indices=tf.expand_dims(jet_idx_capped, axis=-1), updates=py)
jet_pt = tf.math.sqrt(jet_px_new**2 + jet_py_new**2)
return jet_pt
def salt_and_pepper_noise_tf(img_tensor, proportion=0.05):
img_shape = tf.shape(img_tensor, out_type=tf.int64)
float_shape = tf.cast(img_shape, tf.float32)
h, w = img_shape[0], img_shape[1]
fh, fw = float_shape[0], float_shape[1]
num = tf.cast(fh * fw * proportion, tf.int32) // 2 # 添加椒盐噪声的个数
# 椒噪声
hs = tf.random.uniform([num], 0, h - 1, dtype=tf.int64)
ws = tf.random.uniform([num], 0, w - 1, dtype=tf.int64)
hs_ws = tf.stack([hs, ws], axis=1)
noise = tf.zeros(shape=[num, 3], dtype=img_tensor.dtype)
img_tensor = tf.tensor_scatter_nd_add(img_tensor,
indices=hs_ws,
updates=noise)
# 盐噪声
hs = tf.random.uniform([num], 0, h - 1, dtype=tf.int64)
ws = tf.random.uniform([num], 0, w - 1, dtype=tf.int64)
hs_ws = tf.stack([hs, ws], axis=1)
noise = tf.zeros(shape=[num, 3], dtype=img_tensor.dtype) + 255
img_tensor = tf.tensor_scatter_nd_add(img_tensor,
indices=hs_ws,
updates=noise)
return img_tensor
def _backtrack_one_step(flip_matrix, active, curr_job_idx):
"""Take one step in backtracking."""
# Discover the worker that the job originated from, note that this worker
# must exist by construction.
curr_worker_idx = tf.gather(
state["jobs_from_worker"], curr_job_idx, batch_dims=1) - 1
curr_worker_idx = tf.maximum(curr_worker_idx, 0)
update_indices = tf.stack([batch_range, curr_worker_idx, curr_job_idx],
axis=1)
update_indices = tf.maximum(update_indices, 0)
flip_matrix = tf.tensor_scatter_nd_add(flip_matrix, update_indices,
tf.cast(active, tf.int32))
# Discover the (potential) job that the worker originated from.
curr_job_idx = tf.gather(
state["workers_from_job"], curr_worker_idx, batch_dims=1) - 1
# Note that jobs may not be active, and we track that here (before
# adjusting indices so that they are all >= 0 for gather).
active &= curr_job_idx >= 0
curr_job_idx = tf.maximum(curr_job_idx, 0)
update_indices = tf.stack([batch_range, curr_worker_idx, curr_job_idx],
axis=1)
update_indices = tf.maximum(update_indices, 0)
flip_matrix = tf.tensor_scatter_nd_add(flip_matrix, update_indices,
tf.cast(active, tf.int32))
return flip_matrix, active, curr_job_idx
def call(self, obj_vecs, pred_vecs, edges, training=True):
O = obj_vecs.shape[0]
T = pred_vecs.shape[0]
Din, H, Dout = self.input_dim, self.hidden_dim, self.output_dim
# (T, )
s_idx = edges[:, 0]
o_idx = edges[:, 1]
# (T, D)
# cur_s_vecs = obj_vecs[s_idx.numpy()]
# cur_o_vecs = obj_vecs[o_idx.numpy()]
cur_s_vecs = []
cur_o_vecs = []
for idx in s_idx.numpy():
cur_s_vecs.append(obj_vecs[idx])
for idx in o_idx.numpy():
cur_o_vecs.append(obj_vecs[idx])
cur_s_vecs = tf.convert_to_tensor(cur_s_vecs)
cur_o_vecs = tf.convert_to_tensor(cur_o_vecs)
# (T, 3 * D)
cur_t_vecs = tf.concat([cur_s_vecs, pred_vecs, cur_o_vecs], axis=1)
new_t_vecs = self.net1(cur_t_vecs, training=training)
# (T, x)
new_s_vecs = new_t_vecs[:, :H]
new_p_vecs = new_t_vecs[:, H: (H + Dout)]
new_o_vecs = new_t_vecs[:, (H + Dout): (2 * H + Dout)]
# TODO: dtype should be determined by obj_vecs
# (O, H)
pooled_obj_vecs = tf.zeros(shape=(O, H), dtype=tf.float32)
pooled_obj_vecs = tf.tensor_scatter_nd_add(
pooled_obj_vecs, tf.reshape(s_idx, (-1, 1)), new_s_vecs)
pooled_obj_vecs = tf.tensor_scatter_nd_add(
pooled_obj_vecs, tf.reshape(o_idx, (-1, 1)), new_o_vecs)
if self.pooling == 'avg':
# (O, )
obj_counts = tf.zeros(O, dtype=tf.float32)
# (T, )
ones = tf.ones(T, dtype=tf.float32)
obj_counts = tf.tensor_scatter_nd_add(
obj_counts, tf.reshape(s_idx, (-1, 1)), ones)
obj_counts = tf.tensor_scatter_nd_add(
obj_counts, tf.reshape(o_idx, (-1, 1)), ones)
obj_counts = tf.clip_by_value(obj_counts, 1, O)
pooled_obj_vecs = pooled_obj_vecs / tf.reshape(obj_counts, (-1, 1))
new_obj_vecs = self.net2(pooled_obj_vecs, training=training)
return new_obj_vecs, new_p_vecs
def reduce_fn(state, point):
cluster_sums, cluster_weights, num_examples = state
closest_centroid = _find_closest_centroid(centroids, point)
scatter_index = [[closest_centroid]]
cluster_sums = tf.tensor_scatter_nd_add(cluster_sums, scatter_index,
tf.expand_dims(point, axis=0))
cluster_weights = tf.tensor_scatter_nd_add(cluster_weights,
scatter_index, [1])
num_examples += 1
return cluster_sums, cluster_weights, num_examples
def ll_dist(self, gamma, states_t, next_states_t, actions_t, rewards_t, done_mask, weights):
if self.use_double:
states_t = tf.concat([states_t, next_states_t], 0)
# Calculate current state probabilities
net_output = self.net(states_t)
state_action_dist = tf.nn.log_softmax(net_output[:self.batch_size], axis=-1)
state_action_dist = tf.squeeze(tf.gather(state_action_dist, tf.reshape(actions_t, [-1, 1, 1]), batch_dims=1))
# Calculate next state probabilities
target_net_output = tf.nn.softmax(self.tgt_net(next_states_t), axis=-1)
if self.use_double:
next_state_actions = tf.nn.softmax(net_output[self.batch_size:])
next_best_actions = tf.argmax(tf.reduce_sum(next_state_actions, -1), -1)
else:
next_best_actions = tf.argmax(tf.reduce_sum(target_net_output, -1), -1)
next_state_dist = tf.squeeze(tf.gather(target_net_output, tf.reshape(next_best_actions, [-1, 1, 1]),
batch_dims=1))
# Calculate the Bellman operator T to produce Tz
delta_z = (self.v_max - self.v_min) / (self.n_atoms - 1)
support = tf.linspace(self.v_min, self.v_max, self.n_atoms)
Tz = tf.expand_dims(rewards_t, -1) + tf.expand_dims(tf.cast(tf.logical_not(done_mask), tf.float32), -1) * (
gamma ** self.n_steps) * tf.expand_dims(support, 0)
Tz = tf.clip_by_value(Tz, self.v_min, self.v_max)
b = (Tz - self.v_min) / delta_z
l = tf.math.floor(b)
u = tf.math.floor(b)
# Fix disappearing probability mass
eq_mask = tf.equal(l, u)
u_greater = tf.greater(u, 0)
l_less = tf.less(l, self.n_atoms - 1.0)
l = tf.where(tf.logical_and(eq_mask, u_greater), x=l - 1, y=l)
u = tf.where(tf.logical_and(eq_mask, l_less), x=u + 1, y=u)
m = tf.zeros(self.batch_size * self.n_atoms)
offset = tf.linspace(0.0, ((self.batch_size - 1.0) * self.n_atoms), self.batch_size)
offset = tf.reshape(tf.tile(tf.expand_dims(offset, -1), [1, self.n_atoms]), [-1, 1])
m = tf.tensor_scatter_nd_add(
m,
tf.cast(tf.reshape(l, [-1, 1]) + offset, tf.int32),
tf.reshape(next_state_dist * (u - b), [-1])
)
m = tf.tensor_scatter_nd_add(
m,
tf.cast(tf.reshape(u, [-1, 1]) + offset, tf.int32),
tf.reshape(next_state_dist * (b - l), [-1])
)
m = tf.reshape(m, [self.batch_size, self.n_atoms])
m = tf.stop_gradient(m)
# Calculate loss
losses = -tf.reduce_sum(m * state_action_dist, -1)
losses = tf.multiply(weights, losses)
return tf.reduce_mean(losses, -1), losses
def histogram_2d(eta, phi, weights_px, weights_py, eta_range, phi_range, nbins, bin_dtype=tf.float32):
eta_bins = tf.histogram_fixed_width_bins(eta, eta_range, nbins=nbins, dtype=bin_dtype)
phi_bins = tf.histogram_fixed_width_bins(phi, phi_range, nbins=nbins, dtype=bin_dtype)
hist_px = tf.zeros((nbins, nbins), dtype=weights_px.dtype)
hist_py = tf.zeros((nbins, nbins), dtype=weights_py.dtype)
indices = tf.transpose(tf.stack([phi_bins, eta_bins]))
hist_px = tf.tensor_scatter_nd_add(hist_px, indices, weights_px)
hist_py = tf.tensor_scatter_nd_add(hist_py, indices, weights_py)
hist_pt = tf.sqrt(hist_px**2 + hist_py**2)
return hist_pt
def evaluate(self, reward, done, next_observation):
futur_reward = reward
not_done = tf.logical_not(done)
if tf.reduce_any(not_done):
not_done_indicies = tf.where(not_done)
next_value, _ = self.value_policy(next_observation)
try:
tf.tensor_scatter_nd_add(futur_reward, not_done_indicies, self.discount_factor * next_value)
except:
print(futur_reward, not_done_indicies, next_value)
return futur_reward
def parallel_kernel(
prob: Tensor,
gdata: Any,
nnp: Tensor,
kernel_func: Callable[[Any, Tensor, Sequence[int]], Tuple[Tensor, Tensor]],
) -> Tuple[Tensor, Tensor, Tensor]:
"""
kernel for multiprocess to run parallel in DQAS function
:param prob:
:param gdata:
:param nnp:
:param kernel_func:
:return:
"""
sp.random.seed() # make each subprocess run with different random state
# see https://stackoverflow.com/a/6914470/9062180
# it is still not the best way to corporate numpy random and multiprocessing
# see more in https://github.com/numpy/numpy/issues/9650
dtype = tf.float32
p = prob.shape[0]
preset = preset_byprob(prob)
loss, gnnp = kernel_func(gdata, nnp, preset)
gs = tf.tensor_scatter_nd_add(
tf.cast(-prob, dtype=dtype),
tf.constant(list(zip(range(p), preset))),
tf.ones([p], dtype=dtype),
) # \nabla lnp
return loss, gnnp, gs
def kgcnn_ops_tensor_scatter_nd_by_name(segment_name,
tensor,
indices,
updates,
name=None):
"""Scatter operation chosen by name that can replace segment-operations.
Args:
segment_name (str): Operation to update scattered updates. Either 'sum' or 'min' etc.
tensor (tf.tensor): Tensor to scatter updates into.
indices (tf.tensor): Indices to for updates.
updates (tf.tensor): Updates of new entries for tensor.
name (str): Name of the tensor.
Returns:
tf.tensor: Updates scattered into tensor with different update rules.
"""
pool = None
if segment_name in ["segment_mean", "mean", "reduce_mean"]:
pool = tensor_scatter_nd_mean(tensor, indices, updates, name=name)
elif segment_name in ["segment_sum", "sum", "reduce_sum"]:
pool = tf.tensor_scatter_nd_add(tensor, indices, updates, name=name)
elif segment_name in ["segment_max", "max", "reduce_max"]:
pool = tf.tensor_scatter_nd_max(tensor, indices, updates, name=name)
elif segment_name in ["segment_min", "sum", "reduce_min"]:
pool = tf.tensor_scatter_nd_min(tensor, indices, updates, name=name)
else:
raise TypeError("Unknown pooling, choose: 'mean', 'sum', ...")
return pool
def scatter_add_tensor(ref, indices, updates, name=None):
"""
Adds sparse updates to a variable reference.
This operation outputs ref after the update is done. This makes it
easier to chain operations that need to use the reset value.
Duplicate indices: if multiple indices reference the same location,
their contributions add.
Requires updates.shape = indices.shape + ref.shape[1:].
:param ref: A Tensor. Must be one of the following types: float32,
float64, int64, int32, uint8, uint16, int16, int8, complex64, complex128,
qint8, quint8, qint32, half.
:param indices: A Tensor. Must be one of the following types: int32,
int64. A tensor of indices into the first dimension of ref.
:param updates: A Tensor. Must have the same dtype as ref. A tensor of
updated values to add to ref
:param name: A name for the operation (optional).
:return: Same as ref. Returned as a convenience for operations that want
to use the updated values after the update is done.
"""
return tensorflow.tensor_scatter_nd_add(ref, indices, updates, name)
def propagate(self,
in_states,
layer_no,
edge_type_ids,
message_sources,
message_targets,
training,
residuals=None):
# Collect messages across all edge types.
messages = tf.zeros_like(in_states)
for type_index in range(self.num_edge_types):
type_ids = edge_type_ids[type_index]
if tf.shape(type_ids)[0] == 0:
continue
# Retrieve source states and compute type-transformation.
edge_source_states = gather_dense_grad(in_states,
message_sources[type_index])
type_messages = tf.matmul(edge_source_states,
self.type_weights[layer_no][type_index])
if self.add_type_bias:
type_messages += self.type_biases[layer_no][type_index]
messages = tf.tensor_scatter_nd_add(
messages, message_targets[type_index][..., tf.newaxis],
type_messages)
# Concatenate residual messages, if applicable.
if residuals is not None:
messages = tf.concat(residuals + [messages], axis=-1)
# Run GRU for each node.
new_states, _ = self.rnns[layer_no](messages,
in_states,
training=training)
return new_states
def insert_record(histogram, record):
"""Inserts a record to the histogram.
If the record is outside the valid range, it will be dropped.
Args:
histogram: A `tf.Tensor` representing the histogram.
record: A `float` representing the incoming record.
Returns:
A `tf.Tensor` representing updated histgoram with the input record
inserted.
"""
if histogram.shape != (num_bins, ):
raise ValueError(f'Expected shape ({num_bins}, ). '
f'Get {histogram.shape}.')
if record < lower_bound or record >= upper_bound:
return histogram
else:
bin_index = tf.cast(
tf.math.floor((record - lower_bound) / precision), tf.int32)
return tf.tensor_scatter_nd_add(tensor=histogram,
indices=[[bin_index]],
updates=[1])
def linear_small_translation_1d(X,axis,t,name=None):
'''
Translates X along axis by an amount 0<=t<=1. Intuitively, something like
result[...,i...,] approx X[...,i+t,...]
where t is floating point. The shape
of result will be same as the shape of X,
except along axis, where it will reduced by 1.
'''
if name is None:
name='linear_small_translation_1d'
with tf.name_scope(name):
shp=tf.shape(X)
# get two neighbors
stv,szv=construct_1dslice(shp,0,shp[axis]-1,axis)
W = tf.slice(X,stv,szv)
stv=tf.tensor_scatter_nd_add(stv,[[axis]],[1])
E = tf.slice(X,stv,szv)
# linear interpolation does the rest
return linear_interpolation(W,E,t)
def _loss_semantic_segmentation(self, pred_seg, mask_gt, classes, num_obj):
shape_mask = tf.shape(mask_gt)
num_batch = shape_mask[0]
seg_shape = tf.shape(pred_seg)[1]
loss_seg = 0.
for idx in tf.range(num_batch):
seg = pred_seg[idx]
masks = mask_gt[idx]
cls = classes[idx]
objects = num_obj[idx]
# seg shape (69, 69, num_cls)
# resize masks from (100, 138, 138) to (100, 69, 69)
masks = tf.expand_dims(masks, axis=-1)
masks = tf.image.resize(masks, [seg_shape, seg_shape], method=tf.image.ResizeMethod.BILINEAR)
masks = tf.cast((masks > 0.5), seg.dtype)
masks = tf.squeeze(masks)
# obj_mask shape (objects, 138, 138)
obj_mask = masks[:objects]
obj_cls = tf.expand_dims(cls[:objects], axis=-1)
# create empty ground truth (138, 138, num_cls)
seg_gt = tf.zeros_like(seg)
seg_gt = tf.transpose(seg_gt, perm=(2, 0, 1))
seg_gt = tf.tensor_scatter_nd_add(seg_gt, indices=obj_cls, updates=obj_mask)
seg_gt = tf.transpose(seg_gt, perm=(1, 2, 0))
loss_seg += tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(seg_gt, seg))
loss_seg = loss_seg / tf.cast(seg_shape, pred_seg.dtype) ** 2 / tf.cast(num_batch, pred_seg.dtype)
return loss_seg
def _compare_dynamic_gather_nd_with_tf(
test_case, params_shape, static_params_shape, indices_shape
):
params, indices = _random_inputs(params_shape, indices_shape)
i = tf.constant(indices)
with tf.GradientTape() as t:
x = tf.Variable(params)
y = tf.gather_nd(x, i)
dy = t.gradient(y, x)
if isinstance(dy, tf.IndexedSlices):
test_case.assertTrue(
np.array_equal(indices.ravel(), dy.indices.numpy().ravel())
)
zero_params = tf.constant(np.full(params.shape, 0.0, dtype=np.float32))
dy = tf.tensor_scatter_nd_add(zero_params, i, dy.values)
def compare_dy(params_grad):
test_case.assertTrue(np.array_equal(dy.numpy(), params_grad.numpy_list()[0]))
of_y = _of_dynamic_params_gather_nd(
params, indices, static_params_shape, compare_dy
)
test_case.assertTrue(np.array_equal(y.numpy(), of_y))
def body(i, s, S, I, A, R, D, Ru):
U = A + R + D
alpha_t = param_vector[0] + (
param_vector[1] / (tf.constant(1.0) + tf.pow(U, param_vector[2])))
h_1 = (S * I / P) * alpha_t
h_2 = I * param_vector[4]
h_3 = A * param_vector[3]
h_4 = A * param_vector[5]
h_5 = I * param_vector[6] * param_vector[3]
h = tf.stack([h_1, h_2, h_3, h_4, h_5])
Y_store = tf.clip_by_value(
tf.math.floor(tfd.Normal(loc=h, scale=tf.sqrt(h)).sample()), 0.0,
P)
m = tf.matmul(tf.transpose(nu), Y_store)
S = tf.clip_by_value(S + m[0, :], 0.0, P)
I = tf.clip_by_value(I + m[1, :], 0.0, P)
A = tf.clip_by_value(A + m[2, :], 0.0, P)
R = tf.clip_by_value(R + m[3, :], 0.0, P)
D = tf.clip_by_value(D + m[4, :], 0.0, P)
Ru = tf.clip_by_value(Ru + m[5, :], 0.0, P)
s = tf.tensor_scatter_nd_add(s, [[i, 0], [i, 1], [i, 2]],
tf.stack([A, R, D]))
return i + 1, s, S, I, A, R, D, Ru
def layerSave(self,save_mode=True):
'''
For saving intialization values (savemode=True) and then reinstating them without overwriting
the frozen paramerters (savemode=False).
'''
if save_mode:
self.saved_layers = self.net.dense_layers
for i in range(len(self.saved_layers)):
for j in range(len(self.saved_layers[i].trainable_weights)):
params = self.saved_layers[i].trainable_weights[j]
self.saved_layers[i].trainable_weights[j] = params
else:
for i, dlayer in enumerate(self.saved_layers):
for j, layer in enumerate(dlayer.trainable_weights):
# Check for the case of a single-valued layer.
if len(self.frozen_inds[i][j])>0:
tensor=layer
indices=self.frozen_inds[i][j]
updates=self.frozen_vals[i][j]
self.net.dense_layers[i].trainable_weights[j] = tf.tensor_scatter_nd_add(
tensor=tensor,
indices=indices,
updates=updates)
else:
self.net.dense_layers[i].trainable_weights[j] = layer
def graph_sampler(self, batch_size, seed, beta):
#Same as sample method above but specialised for graph compilation
sample = tf.zeros([batch_size, self.L, self.L, 1], tf.float32)
tf_binomial = tf.random.stateless_binomial
full_ones = tf.ones([batch_size], tf.int32)
full_zeros = tf.zeros_like(full_ones)
r = self.learn_range
for i in range(self.L):
for j in range(self.L):
seed.assign((seed * 1664525 + 1013904223) % 2**31)
sub_sample = sample[:,
np.maximum(i - 1, 0):i + 1,
np.maximum(j - r, 0):np.
minimum(j + r + 1, self.L)]
x_hat = self.call(sub_sample, beta)
i_h = tfm.minimum(i, 1)
j_h = tfm.minimum(j, r)
probs = 0.5 if i == 0 and j == 0 else x_hat[:, i_h, j_h, 0]
indices = tf.stack([
tf.range(batch_size), i * full_ones, j * full_ones,
full_zeros
], 1)
updates = tf_binomial([batch_size], seed, 1., probs,
tf.float32) * 2 - 1
sample = tf.tensor_scatter_nd_add(sample,
tf.cast(indices, tf.int32),
updates)
#x_hat = self.call(sample)
if self.z2:
seed.assign((seed * 1664525 + 1013904223) % 2**31)
flip = tf_binomial([batch_size, 1, 1, 1], seed, 1., 0.5,
tf.float32) * 2 - 1
sample = sample * flip
return sample
def _get_labels_embed(
label_ids: tf.Tensor, all_labels_embed: tf.Tensor
) -> tf.Tensor:
# instead of processing labels again, gather embeddings from
# all_labels_embed using label ids
indices = tf.cast(label_ids[:, :, 0], tf.int32)
# Find padding indices. They should have a value equal to `LABEL_PAD_ID`
padding_indices = tf.where(tf.equal(indices, LABEL_PAD_ID))
# Create a tensor of values with sign opposite to `LABEL_PAD_ID` which
# will serve as updates to original `indices`
updates_to_indices = (
tf.ones((tf.shape(padding_indices)[0]), dtype=tf.int32) * -1 * LABEL_PAD_ID
)
# Add the updates tensor to indices with padding.
# So, effectively all indices with `LABEL_PAD_ID=-1`
# become 0 because updates contain 1s.
# This is fine because we don't change the original non-padding label
# indices but only make the padding indices 'compatible'
# for the `tf.gather` op below.
indices_to_gather = tf.cast(
tf.tensor_scatter_nd_add(indices, padding_indices, updates_to_indices),
tf.int32,
)
labels_embed = tf.gather(all_labels_embed, indices_to_gather)
return labels_embed
def _tf_scatter_add(tensor, indices, updates):
# manual implementation of pytorch._scatter_add(dim=2, src=tensor, idx=indeces, out=updaties)
original_tensor = tensor
# expand index value from vocab size
indices = tf.reshape(indices, shape=[-1, tf.shape(indices)[-1]])
indices_add = tf.expand_dims(tf.range(0,
tf.shape(indices)[0], 1) *
(tf.shape(tensor)[-1]),
axis=-1)
indices += indices_add
# resize
tensor = tf.reshape(tensor, shape=[-1])
indices = tf.reshape(indices, shape=[-1, 1])
updates = tf.reshape(updates, shape=[-1])
#check_
"""
update = tensor.shape[indices.shape[-1]:]
res = indices.shape[:-1] + update
"""
#same Torch scatter_add_
#scatter = tf.scatter_nd_add(tensor, indices, updates)
scatter = tf.tensor_scatter_nd_add(tensor, indices, updates)
scatter = tf.reshape(scatter,
shape=[
tf.shape(original_tensor)[0],
tf.shape(original_tensor)[1], -1
])
return scatter
def body(i, s, S, I, A, R, D, Ru):
"""Single update for one day."""
U = A / (tf.pow(10.0, param_vector[8]))
alpha_t = param_vector[0] + (
param_vector[1] / (tf.constant(1.0) + tf.pow(U, param_vector[6])))
h_1 = (S * I / P) * alpha_t
h_2 = I * param_vector[3]
h_3 = A * param_vector[2]
h_4 = A * param_vector[4]
h_5 = I * param_vector[2] * param_vector[5]
h = tf.stack([h_1, h_2, h_3, h_4, h_5])
normal_sample = tfd.Normal(loc=h, scale=tf.sqrt(h)).sample()
Y_store = tf.clip_by_value(tf.math.floor(normal_sample), 0.0, P)
m = tf.matmul(tf.transpose(MIXING_MATRIX), Y_store)
# Note: Simple vectorisation suppresses parameter update in loop.
S = tf.clip_by_value(S + m[0, :], 0.0, P)
I = tf.clip_by_value(I + m[1, :], 0.0, P)
A = tf.clip_by_value(A + m[2, :], 0.0, P)
R = tf.clip_by_value(R + m[3, :], 0.0, P)
D = tf.clip_by_value(D + m[4, :], 0.0, P)
Ru = tf.clip_by_value(Ru + m[5, :], 0.0, P)
s = tf.tensor_scatter_nd_add(tensor=s,
indices=[[i, 0], [i, 1], [i, 2]],
updates=tf.stack([A, R, D]))
return i+1, s, S, I, A, R, D, Ru
def accumulate_client_votes(vote_accumulator, example):
"""Accumulates client votes on prefix extensions."""
example = tf.strings.lower(example)
# Append the default terminator to the example.
example = tf.strings.join([example, default_terminator])
# Compute effective round number.
effective_round_num = tf.math.floordiv(round_num, num_sub_rounds)
if tf.strings.length(example) < effective_round_num:
return vote_accumulator
else:
discovered_prefixes_index = discovered_prefixes_table.lookup(
tf.strings.substr(example, 0, effective_round_num))
possible_prefix_extensions_index = possible_prefix_extensions_table.lookup(
tf.strings.substr(example, effective_round_num, 1))
# If the character extension is not in the alphabet, or the prefix is not
# already in the discovered prefixes, do not add client's vote.
if tf.math.logical_or(
tf.math.equal(possible_prefix_extensions_index,
tf.constant(DEFAULT_VALUE)),
tf.math.equal(discovered_prefixes_index,
tf.constant(DEFAULT_VALUE))):
return vote_accumulator
else:
indices = [[
discovered_prefixes_index, possible_prefix_extensions_index
]]
updates = tf.constant([1])
return tf.tensor_scatter_nd_add(vote_accumulator, indices,
updates)
def apply_scatter_nd_add(tensor, updates, indices, tf_int, tf_float):
""" applies the tensor_scatter_nd_add over the batch dimension
"""
out = Lambda(lambda entry: K.map_fn(
lambda entry: tf.tensor_scatter_nd_add(entry[0], entry[1], entry[2]),
entry,
dtype=tf_float))([tensor, indices, updates])
return out
def _compare_tensor_scatter_nd_add_with_tf(test_case, params_shape,
indices_shape, updates_shape,
device_type, mirrored):
params, updates, indices = _random_inputs(params_shape, indices_shape,
updates_shape, True)
params_const = tf.constant(params)
indices_const = tf.constant(indices)
updates_const = tf.constant(updates)
with tf.GradientTape() as t1:
params_var = tf.Variable(params)
tf_out1 = tf.tensor_scatter_nd_add(params_var, indices_const,
updates_const)
tf_params_grad = t1.gradient(tf_out1, params_var)
with tf.GradientTape() as t2:
updates_var = tf.Variable(updates)
tf_out2 = tf.tensor_scatter_nd_add(params_const, indices_const,
updates_var)
tf_updates_grad = t2.gradient(tf_out2, updates_var)
test_case.assertTrue(np.allclose(tf_out1.numpy(), tf_out2.numpy()))
def compare_params_grad(of_params_grad):
tf_params_grad_np = tf_params_grad.numpy()
of_params_grad_np = (of_params_grad.numpy_list()[0]
if mirrored else of_params_grad.numpy())
test_case.assertTrue(np.allclose(tf_params_grad_np, of_params_grad_np))
def compare_updates_grad(of_updates_grad):
tf_updates_grad_np = tf_updates_grad.numpy()
of_updates_grad_np = (of_updates_grad.numpy_list()[0]
if mirrored else of_updates_grad.numpy())
test_case.assertTrue(
np.allclose(tf_updates_grad_np, of_updates_grad_np))
of_out = _of_tensor_scatter_nd_add(
params,
indices,
updates,
device_type,
mirrored,
compare_params_grad,
compare_updates_grad,
)
test_case.assertTrue(np.allclose(tf_out1.numpy(), of_out))
def _scatter_element_add_tf(tensor, index, value):
"""In-place addition of a multidimensional value over various
indices of a tensor."""
import tensorflow as tf
indices = tf.expand_dims(index, 0)
value = tf.cast(tf.expand_dims(value, 0), tensor.dtype)
return tf.tensor_scatter_nd_add(tensor, indices, value)
def scatter_sum(src, index, dim: int = -1, out=None, dim_size=None):
index = broadcast(index, src, dim)
if out is None:
size = src.get_shape().as_list()
if dim_size is not None:
size[dim] = dim_size
elif index.numel() == 0:
size[dim] = 0
else:
size[dim] = int(index.max()) + 1
out = tf.zeros(size, dtype=src.dtype)
print(out)
print(index)
print(src)
return tf.tensor_scatter_nd_add(out, index, src)
else:
return tf.tensor_scatter_nd_add(out, index, src)
def lazy_marginal(self, instruments, numeraire, time, book):
batch_size = instruments.shape[0]
book_size = book.book_size
timesteps = len(time) - 1
marginal_prices = tf.zeros((batch_size, book_size, timesteps + 1),
FLOAT_DTYPE)
marginal_deltas = tf.zeros((batch_size, book_size, timesteps + 1),
FLOAT_DTYPE)
grid = tf.constant(list(itertools.product(
range(batch_size), range(book_size), range(timesteps + 1))))
for book_idx in tf.range(book_size):
entry = book.derivatives[book_idx]
instrument = instruments[:, entry["link"], :]
sign = entry["exposure"]
value = sign * entry["derivative"].value(
time, instrument, numeraire)
delta = sign * entry["derivative"].delta(
time, instrument, numeraire)
indices = tf.boolean_mask(grid, grid[:, 1] == book_idx)
marginal_prices = tf.tensor_scatter_nd_add(
marginal_prices, indices, tf.reshape(value, -1))
marginal_deltas = tf.tensor_scatter_nd_add(
marginal_deltas, indices, tf.reshape(delta, -1))
prices = tf.reduce_sum(marginal_prices, axis=1)
deltas = tf.zeros((batch_size, book.instrument_dim, timesteps + 1))
links = [entry["link"] for entry in book.derivatives]
for path_idx in range(batch_size):
for time_idx in range(timesteps + 1):
indices = tf.stack([
path_idx * tf.ones(book.instrument_dim, INT_DTYPE),
tf.range(book.instrument_dim),
time_idx * tf.ones(book.instrument_dim, INT_DTYPE)], 1)
updates = tf.math.bincount(
links,
marginal_deltas[path_idx, :, time_idx],
book.instrument_dim)
deltas = tf.tensor_scatter_nd_add(deltas, indices, updates)
return prices, deltas
def call(self, inputs):
# NOTE: it is for now impossible to do slice assignment in tensorflow
# there are some on-going GH issues or SO questions but for now
# tensor_scatter_nd_add seems to be the only way to go.
# https://stackoverflow.com/questions/62092147/how-to-efficiently-assign-to-a-slice-of-a-tensor-in-tensorflow
in_shape = tf.shape(inputs)
batch_size = in_shape[0]
height = in_shape[1]
width = in_shape[2]
# the number of channels can't be unknown for the convolutions
n_channels = inputs.shape[3] // 4
outputs = tf.zeros([batch_size, 2 * height, 2 * width, n_channels])
# for now we only consider greyscale
x1 = inputs[..., 0:n_channels] / 2
x2 = inputs[..., n_channels:2 * n_channels] / 2
x3 = inputs[..., 2 * n_channels:3 * n_channels] / 2
x4 = inputs[..., 3 * n_channels:4 * n_channels] / 2
# in the following, E denotes even and O denotes odd
x_EE = x1 - x2 - x3 + x4
x_OE = x1 - x2 + x3 - x4
x_EO = x1 + x2 - x3 - x4
x_OO = x1 + x2 + x3 + x4
# now the preparation to tensor_scatter_nd_add
height_range_E = 2 * tf.range(height)
height_range_O = height_range_E + 1
width_range_E = 2 * tf.range(width)
width_range_O = width_range_E + 1
# this transpose allows to only index the varying dimensions
# only the first dimensions can be indexed in tensor_scatter_nd_add
# we also need to match the indices with the updates reshaping
scatter_nd_perm = [2, 1, 3, 0]
outputs_reshaped = tf.transpose(outputs, perm=scatter_nd_perm)
combos_list = [
((height_range_E, width_range_E), x_EE),
((height_range_O, width_range_E), x_OE),
((height_range_E, width_range_O), x_EO),
((height_range_O, width_range_O), x_OO),
]
for (height_range, width_range), x_comb in combos_list:
h_range, w_range = tf.meshgrid(height_range, width_range)
h_range = tf.reshape(h_range, (-1, ))
w_range = tf.reshape(w_range, (-1, ))
combo_indices = tf.stack([w_range, h_range], axis=-1)
combo_reshaped = tf.transpose(x_comb, perm=scatter_nd_perm)
outputs_reshaped = tf.tensor_scatter_nd_add(
outputs_reshaped,
indices=combo_indices,
updates=tf.reshape(combo_reshaped,
(-1, n_channels, batch_size)),
)
inverse_scatter_nd_perm = [3, 1, 0, 2]
outputs = tf.transpose(outputs_reshaped, perm=inverse_scatter_nd_perm)
return outputs
