diff --git a/ads/catalog/notebook.py b/ads/catalog/notebook.py
index c8ac60850..5af237d0c 100644
--- a/ads/catalog/notebook.py
+++ b/ads/catalog/notebook.py
@@ -22,6 +22,8 @@
     CreateNotebookSessionDetails,
     NotebookSession,
     NotebookSessionConfigurationDetails,
+    NotebookSessionConfigDetails,
+    NotebookSessionShapeConfigDetails,
 )
 from oci.exceptions import ServiceError
 from types import MethodType
@@ -295,6 +297,9 @@ def create_notebook_session(
         shape=None,
         block_storage_size_in_gbs=None,
         subnet_id=None,
+        ocpus=None,
+        memory_in_gbs=None,
+        private_endpoint_id=None,
         **kwargs,
     ):
         """
@@ -310,11 +315,15 @@ def create_notebook_session(
             The value to assign to the shape property of this NotebookSessionConfigurationDetails.
             Allowed values for this property are: "VM.Standard.E2.2", "VM.Standard.E2.4",
             "VM.Standard.E2.8", "VM.Standard2.1", "VM.Standard2.2", "VM.Standard2.4", "VM.Standard2.8",
-            "VM.Standard2.16","VM.Standard2.24".
+            "VM.Standard2.16","VM.Standard2.24". Flexible shapes also supported but require addition of ocpus and memory_in_gbs parameters
         block_storage_size_in_gbs: int, required
             Size of the block storage drive. Limited to values between 50 (GB) and 1024 (1024GB = 1TB)
-        subnet_id: str, required
-            The OCID of the subnet resource where the notebook is to be created.
+        subnet_id: str, optional
+            The OCID of the subnet resource where the notebook is to be created. If no subnet_id is specified the session will use managed egress
+        ocpus: int, optional
+            OCPUs assigned to flexible shapes, such as VM.Standard.E4.Flex
+        memory_in_gbs: int, optional
+            memory (in GBs) assigned to flexible shapes, such as VM.Standard.E4.Flex
         kwargs: dict, optional
             Additional kwargs passed to `DataScienceClient.create_notebook_session()`
 
@@ -326,18 +335,34 @@ def create_notebook_session(
         ------
         KeyError: If the resource was not found or do not have authorization to access that resource.
         """
-        notebook_session_configuration_details = NotebookSessionConfigurationDetails(
-            shape=shape,
-            block_storage_size_in_gbs=block_storage_size_in_gbs,
-            subnet_id=subnet_id,
-        )
+
+        # build configuration kwargs
+        config_kwargs = dict(shape=shape,)
+
         project_id = PROJECT_OCID if project_id is None else project_id
+
+        if block_storage_size_in_gbs:
+            config_kwargs["block_storage_size_in_gbs"] = block_storage_size_in_gbs
+
+        
+        if subnet_id:
+            config_kwargs["subnet_id"] = subnet_id
+        if private_endpoint_id:
+            config_kwargs["private_endpoint_id"] = private_endpoint_id
+        if ocpus is not None and memory_in_gbs is not None:
+            config_kwargs["notebook_session_shape_config_details"] = NotebookSessionShapeConfigDetails(ocpus=ocpus, memory_in_gbs=memory_in_gbs)
+
+        # notebook config object
+        notebook_cfg = NotebookSessionConfigDetails(**config_kwargs)
+
+        # create request details
         create_notebook_details = CreateNotebookSessionDetails(
             display_name=display_name,
             project_id=project_id,
             compartment_id=self.compartment_id,
-            notebook_session_configuration_details=notebook_session_configuration_details,
+            notebook_session_config_details=notebook_cfg,  
         )
+
         try:
             create_notebook_response = self.ds_client.create_notebook_session(
                 create_notebook_details, **kwargs