tf.compat.v1.variable_axis_size_partitioner