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

cos_y = tf.math.cos(ang_y)
            sin_y = tf.math.sin(ang_y)
            # Do the rotations
            for j in range(part_num + 1):
                px = kin.x_component(generated_particles[j])
                py = kin.y_component(generated_particles[j])
                # Rotate about z
                # TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it?
                generated_particles[j] = tf.concat([cos_z * px - sin_z * py,
                                                    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

cos_z = (tf.constant(2.0, dtype=tf.float64) * tf.random.uniform((n_events, 1), dtype=tf.float64)
                         - tf.constant(1.0, dtype=tf.float64))
                sin_z = tf.sqrt(tf.constant(1.0, dtype=tf.float64) - cos_z * cos_z)
                ang_y = (tf.constant(2.0, dtype=tf.float64) * tf.constant(pi, dtype=tf.float64)
                         * tf.random.uniform((n_events, 1), dtype=tf.float64))
            cos_y = tf.math.cos(ang_y)
            sin_y = tf.math.sin(ang_y)
            # Do the rotations
            for j in range(part_num + 1):
                px = kin.x_component(generated_particles[j])
                py = kin.y_component(generated_particles[j])
                # Rotate about z
                # TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it?
                generated_particles[j] = tf.concat([cos_z * px - sin_z * py,
                                                    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,