How to use the phasespace.phasespace function in phasespace

def run_test(n_particles, test_prefix):
    sess = tf.Session()
    first_run_n_events = 100
    main_run_n_events = 100000
    n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64, use_resource=True)
    sess.run(n_events.initializer)

    generate = phasespace.nbody_decay(decays.B0_MASS,
                                      [decays.PION_MASS] * n_particles) \
        .generate_tensor(n_events)
    weights1, _ = sess.run(generate)  # only generate to test change in n_events
    assert len(weights1) == first_run_n_events

    # change n_events and run again
    n_events.load(main_run_n_events, session=sess)
    weights, particles = sess.run(generate)
    parts = np.concatenate([particles[f"p_{part_num}"] for part_num in range(n_particles)], axis=1)
    histos = [make_norm_histo(parts[:, coord],
                              range_=(-3000 if coord % 4 != 3 else 0, 3000),
                              weights=weights)
              for coord in range(parts.shape[1])]
    weight_histos = make_norm_histo(weights, range_=(0, 1+1e-8))
    ref_histos, ref_weights = create_ref_histos(n_particles)
    p_values = np.array([ks_2samp(histos[coord], ref_histos[coord])[1]