@@ -15,31 +15,49 @@
 # limitations under the License.
 
 from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema import predict
-from google.cloud.aiplatform.v1beta1.schema import trainingjob
+from google.cloud.aiplatform.v1.schema import predict
+from google.cloud.aiplatform.v1.schema import trainingjob
+from google.cloud.aiplatform.v1beta1.schema import predict as predict_v1beta1
+from google.cloud.aiplatform.v1beta1.schema import predict as trainingjob_v1beta1
+
+# import the v1 submodules for enhancement
+from google.cloud.aiplatform.v1.schema.predict.instance_v1 import types as instance
+from google.cloud.aiplatform.v1.schema.predict.params_v1 import types as params
+from google.cloud.aiplatform.v1.schema.predict.prediction_v1 import types as prediction
+from google.cloud.aiplatform.v1.schema.trainingjob.definition_v1 import (
+types as definition,
+)
+
+# import the v1beta1 submodules for enhancement
 from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1 import (
-types as instance,
+types as instance_v1beta1,
 )
 from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1 import (
-types as params,
+types as params_v1beta1,
 )
 from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1 import (
-types as prediction,
+types as prediction_v1beta1,
 )
 from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1 import (
-types as definition,
+types as definition_v1beta1,
 )
 
 __all__ = (
 "predict",
 "trainingjob",
+"predict_v1beta1",
+"trainingjob_v1beta1",
 )
 
 enhanced_types_packages = [
 instance,
 params,
 prediction,
 definition,
+instance_v1beta1,
+params_v1beta1,
+prediction_v1beta1,
+definition_v1beta1,
 ]
 
 for pkg in enhanced_types_packages:

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.aiplatform.v1.schema import predict
+from google.cloud.aiplatform.v1.schema import trainingjob
+
+
+__all__ = (
+"predict",
+"trainingjob",
+)

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.aiplatform.v1.schema.predict import instance
+from google.cloud.aiplatform.v1.schema.predict import params
+from google.cloud.aiplatform.v1.schema.predict import prediction
+
+__all__ = (
+"instance",
+"params",
+"prediction",
+)

@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.aiplatform.v1.schema.trainingjob import definition
+
+__all__ = ("definition",)

@@ -13,7 +13,10 @@
 # limitations under the License.
 from __future__ import absolute_import
 
-from google.cloud.aiplatform.v1beta1.schema.trainingjob import definition
+from google.cloud.aiplatform.v1.schema.trainingjob import definition
+from google.cloud.aiplatform.v1beta1.schema.trainingjob import (
+definition as definition_v1beta1,
+)
 
 ModelType = definition.AutoMlImageClassificationInputs().ModelType
 test_training_input = definition.AutoMlImageClassificationInputs(
@@ -23,14 +26,36 @@
 disable_early_stopping=False,
 )
 
+ModelType_v1beta1 = definition_v1beta1.AutoMlImageClassificationInputs().ModelType
+test_training_input_v1beta1 = definition.AutoMlImageClassificationInputs(
+multi_label=True,
+model_type=ModelType_v1beta1.CLOUD,
+budget_milli_node_hours=8000,
+disable_early_stopping=False,
+)
+
 
-def test_exposes_to_value_method():
+# Test the v1 enhanced types.
+def test_exposes_to_value_method_v1():
 assert hasattr(test_training_input, "to_value")
 
 
-def test_exposes_from_value_method():
+def test_exposes_from_value_method_v1():
 assert hasattr(test_training_input, "from_value")
 
 
-def test_exposes_from_map_method():
+def test_exposes_from_map_method_v1():
 assert hasattr(test_training_input, "from_map")
+
+
+# Test the v1beta1 enhanced types.
+def test_exposes_to_value_method_v1beta1():
+assert hasattr(test_training_input_v1beta1, "to_value")
+
+
+def test_exposes_from_value_method_v1beta1():
+assert hasattr(test_training_input_v1beta1, "from_value")
+
+
+def test_exposes_from_map_method_v1beta1():
+assert hasattr(test_training_input_v1beta1, "from_map")