How to use the edward2.make_log_joint_fn function in edward2

def edward_inference(data, args):
    features = tf.cast(data[0], dtype=tf.float32)
    labels = tf.cast(data[1], dtype=tf.int32)

    tf.enable_v2_behavior()
    print("GPU(s) available", tf.test.is_gpu_available())

    log_joint = ed.make_log_joint_fn(edward_model)
    @tf.function  # use graph mode
    def target_log_prob_fn(coeffs):
        return log_joint(features=features, coeffs=coeffs, labels=labels)

    step_size = 0.0015
    kernel = edward2_nuts.kernel
    coeffs_samples = []
    target_log_prob = None
    grads_target_log_prob = None
    seed_stream = tfp.distributions.SeedStream(args.seed, "main")
    coeffs = tf.random.uniform(shape=features.shape[1:],
                               minval=-2,
                               maxval=2,
                               dtype=features.dtype,
                               seed=seed_stream())