context_features = {
        "a":
            parsing_ops.RaggedFeature(
                value_key="a_values",
                partitions=[