How to use the phasespace.kinematics.boost_components function in phasespace

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 = []