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

kin.y_component(generated_particles[j]),
                                                    sin_y * px + cos_y * pz,
                                                    kin.time_component(generated_particles[j])],
                                                   axis=1)
            if part_num == (n_particles - 1):
                break
            betas = (pds[part_num] / tf.sqrt(tf.square(pds[part_num]) + tf.square(inv_masses[part_num])))
            generated_particles = [kin.lorentz_boost(part,
                                                     tf.concat([zero_component,
                                                                betas,
                                                                zero_component],
                                                               axis=1))
                                   for part in generated_particles]
            part_num += 1
        # Final boost of all particles
        generated_particles = [kin.lorentz_boost(part, p_top_boost)
                               for part in generated_particles]
        return tf.reshape(weights, (n_events,)), tf.reshape(w_max, (n_events,)), generated_particles, masses

sin_z * px + cos_z * py,
                                                    kin.z_component(generated_particles[j]),
                                                    kin.time_component(generated_particles[j])],
                                                   axis=1)
                # Rotate about y
                px = kin.x_component(generated_particles[j])
                pz = kin.z_component(generated_particles[j])
                generated_particles[j] = tf.concat([cos_y * px - sin_y * pz,
                                                    kin.y_component(generated_particles[j]),
                                                    sin_y * px + cos_y * pz,
                                                    kin.time_component(generated_particles[j])],
                                                   axis=1)
            if part_num == (n_particles - 1):
                break
            betas = (pds[part_num] / tf.sqrt(tf.square(pds[part_num]) + tf.square(inv_masses[part_num])))
            generated_particles = [kin.lorentz_boost(part,
                                                     tf.concat([zero_component,
                                                                betas,
                                                                zero_component],
                                                               axis=1))
                                   for part in generated_particles]
            part_num += 1
        # Final boost of all particles
        generated_particles = [kin.lorentz_boost(part, p_top_boost)
                               for part in generated_particles]
        return tf.reshape(weights, (n_events,)), tf.reshape(w_max, (n_events,)), generated_particles, masses