How to use the polyaxon.layers.Dense function in polyaxon

def test_variable_sharing(self):
        l = plx.layers.Dense(units=1)
        x = tf.placeholder(dtype=tf.float32, shape=[1, 1])
        y = tf.placeholder(dtype=tf.float32, shape=[2, 1])

        lx = l(x)
        ly = l(y)

        init_all_op = tf.global_variables_initializer()
        assign_op = l.variables[0].assign_add([[1]])

        with self.test_session() as sess:
            sess.run(init_all_op)
            lx_results = lx.eval({x: [[1]]})
            ly_results = ly.eval({y: [[1], [1]]})
            assert len(lx_results) == 1
            assert len(ly_results) == 2

def graph_fn1(mode, x):
            return plx.layers.Dense(units=1)(x)

def graph_fn2(mode, x):
            return plx.layers.Dense(units=1, trainable=False)(x)

def graph_fn1(mode, x):
            return plx.layers.Dense(units=1)(x)

def encoder_fn(mode, features):
    x = plx.layers.Dense(units=128)(features)
    return plx.layers.Dense(units=256)(x)

def graph_fn(mode, features):
    x = plx.layers.Conv2D(filters=32, kernel_size=3, strides=1, activation='elu',
                          kernel_regularizer=l2(0.01))(features['image'])
    x = plx.layers.MaxPooling2D(pool_size=2)(x)
    x = plx.layers.Conv2D(filters=64, kernel_size=3, activation='relu',
                          kernel_regularizer=l2(0.01))(x)
    x = plx.layers.MaxPooling2D(pool_size=2)(x)
    x = plx.layers.Flatten()(x)
    x = plx.layers.Dense(units=128, activation='tanh')(x)
    x = plx.layers.Dropout(rate=0.8)(x)
    x = plx.layers.Dense(units=256, activation='tanh')(x)
    x = plx.layers.Dropout(rate=0.8)(x)
    x = plx.layers.Dense(units=10)(x)
    return x

def graph_fn(mode, features):
        x = plx.layers.Dense(units=32, activation='tanh')(features['X'])
        return plx.layers.Dense(units=1, activation='sigmoid')(x)