Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
min_mass = tf.broadcast_to(recurse_stable(child), (n_events, 1))
mass = child.get_mass(min_mass, max_mass, n_events)
max_mass -= mass
masses.append(tf.reshape(mass, (n_events, 1)))
masses = tf.concat(masses, axis=-1)
# if masses.shape.ndims == 1:
# masses = tf.expand_dims(masses, axis=0)
available_mass = top_mass - tf.reduce_sum(masses, axis=1, keepdims=True)
mass_check = tf.assert_greater_equal(available_mass, tf.zeros_like(available_mass, dtype=tf.float64),
message="Forbidden decay",
name="mass_check")
with tf.control_dependencies([mass_check]):
available_mass = tf.identity(available_mass)
# Calculate the max weight, initial beta, etc
w_max = self._get_w_max(available_mass, masses)
p_top_boost = kin.boost_components(p_top)
# Start the generation
random_numbers = tf.random.uniform((n_events, n_particles - 2), dtype=tf.float64)
random = tf.concat([tf.zeros((n_events, 1), dtype=tf.float64),
tf.sort(random_numbers, axis=1),
tf.ones((n_events, 1), dtype=tf.float64)],
axis=1)
if random.shape[1].value is None:
random.set_shape((None, n_particles))
# random = tf.expand_dims(random, axis=-1)
sum_ = tf.zeros((n_events, 1), dtype=tf.float64)
inv_masses = []
# TODO(Mayou36): rewrite with cumsum?
for i in range(n_particles):
sum_ += tf.gather(masses, [i], axis=1)
inv_masses.append(tf.gather(random, [i], axis=1) * available_mass + sum_)
pds = []