Update docs and add examples for Compute Plan and In-memory model usa…

…ge. (#2407)

* Update docs and add examples for Compute Plan and In-memory model usage

* Address comments

coremltools/models/compute_device.py

@@ -50,14 +50,17 @@ def get_all_compute_devices(
 
         Returns
         -------
-        The accessible compute devices.
+        List[MLComputeDevice]
+            The accessible compute devices.
 
         Examples
         --------
         .. sourcecode:: python
+
             compute_devices = (
                 coremltools.models.compute_device.MLComputeDevice.get_all_compute_devices()
             )
+        
         """
         return _MLModelProxy.get_all_compute_devices()
 
@@ -105,7 +108,7 @@ def total_core_count(self) -> int:
 
         Returns
         -------
-        int:
+        int
             The total number of cores in the Neural Engine.
 
         Examples

coremltools/models/compute_plan.py

@@ -258,11 +258,13 @@ def load_from_path(cls, compiled_model_path: str) -> "MLModelStructure":
 
         Returns
         -------
-        MLModelStructure: An instance of MLModelStructure.
+        MLModelStructure
+            An instance of MLModelStructure.
 
         Examples
         --------
         .. sourcecode:: python
+        
             model_structure = coremltools.models.compute_plan.MLModelStructure.load_from_path(
                 model.get_compiled_path()
             )
@@ -279,6 +281,7 @@ def load_from_path(cls, compiled_model_path: str) -> "MLModelStructure":
             else:
                 # The model type is something else.
                 pass
+        
         """
 
         if _MLModelProxy is None:
@@ -371,7 +374,7 @@ def get_compute_device_usage_for_neuralnetwork_layer(
 
         Returns
         -------
-        Optional[MLComputePlanDeviceUsage]:
+        Optional[MLComputePlanDeviceUsage]
             The anticipated compute devices that would be used for executing the layer or ``None`` if the usage couldn't be determined.
         """
         return self.__proxy__.get_compute_device_usage_for_neuralnetwork_layer(layer)
@@ -418,6 +421,7 @@ def load_from_path(
         Examples
         --------
         .. sourcecode:: python
+
             compute_plan = coremltools.models.compute_plan.MLComputePlan.load_from_path(
                 model.get_compiled_path()
             )

coremltools/models/model.py

@@ -216,7 +216,7 @@ def from_path(
             The file path to the compiled model.
 
         Returns
-        ----------
+        -------
         MLModelAsset
             An instance of MLModelAsset created from the specified path.
         """
@@ -240,7 +240,7 @@ def from_memory(
             A dictionary with blob path as the key and blob data as the value.
 
         Returns
-        ----------
+        -------
         MLModelAsset
             An instance of MLModelAsset created from the provided memory data.
         """

docs-guides/conf.py

@@ -7,9 +7,9 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
 project = 'Core ML Tools Guide'
-copyright = '2023, Apple Inc'
+copyright = '2024, Apple Inc'
 author = 'Apple'
-release = '7.0'
+release = '8.1'
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

docs-guides/source/mlmodel-utilities.md

@@ -257,3 +257,85 @@ Optional arguments:
 Special values for `input_names` and `output_names` arguments:
 * an empty list means nothing will be modified (default for `input_names`)
 * a list containing `"*"` string means all relevant inputs/outputs will be modified (those that will match the `from_type` type)
+
+## Compute Plan
+
+In certain situations, you may want to evaluate the computational needs of a Core ML model before deploying it. 
+The `MLComputePlan` class is designed for this purpose, allowing you to get insights into the resources and costs
+associated with using the model.
+
+Here’s what you can do with `MLComputePlan`:
+- Model Structure: Examine the model structure.
+- Compute Device Usage: Get insights into the compute devices that would be used for executing an ML Program operation/ NeuralNetwork layer.
+- Estimated Cost: Get the estimated cost of executing an ML Program operation.
+
+An example on how to use `MLComputePlan` to get the estimated cost and compute device usages for the operations in an ML Program:
+
+```python
+import coremltools as ct
+# Path to the compiled ML Program model.
+compiled_model_path = "my_model.mlmodelc"
+# Load the compute plan of a model.
+compute_plan = ct.models.MLComputePlan.compute_plan.load_from_path(
+    path=compiled_model_path,
+    compute_units=ct.ComputeUnits.ALL,
+)
+# Get the model structure.
+program = compute_plan.model_structure.program
+mainFunction = program.functions["main"]
+for operation in mainFunction.block.operations:
+    # Get the compute device usage for the operation.
+    compute_device_usage = (
+        compute_plan.get_compute_device_usage_for_mlprogram_operation(operation)
+    )
+    # Get the estimated cost of executing the operation.
+    estimated_cost = compute_plan.get_estimated_cost_for_mlprogram_operation(operation)
+```
+
+## In-memory Model
+If you are using an in-memory model in your application, you can easily test the workflow with `MLModelAsset`. The `MLModelAsset` class includes 
+the `MLModelAsset.from_memory` API, which enables you to load a model directly from the model's in-memory specification data. Once loaded, you
+can use the model to make predictions.
+
+An example on how to use `MLModelAsset` to load an `MLCompiledModel` from in-memory specification data:
+
+```python
+import coremltools as ct
+# Path to the model.
+model = MLModel("my_model.model")
+model_spec = model.get_spec()
+spec_data = model_spec.SerializeToString()
+asset = ct.models.model.MLModelAsset.from_memory(spec_data=spec_data)
+compiled_model = ct.models.CompiledMLModel.from_asset(asset=asset)
+result = compiled_model.predict(
+    {
+        "x": np.array([1.0]),
+        "y": np.array([2.0]),
+    }
+)
+```
+
+Another example on how to use `MLModelAsset` to load a MLCompiledModel from in-memory specification data where the specification has external blob file references :
+
+
+```python
+import coremltools as ct
+# Path to the model.
+mlmodel = MLModel("my_model.mlpackage")
+weight_file_path = mlmodel.weights_dir + "/weight.bin"
+with open(weight_file_path, "rb") as file:
+    weights_data = file.read()
+    model_spec = model.get_spec()
+    spec_data = model_spec.SerializeToString()
+    # Provide the weights data as `blob_mapping`.
+    asset = ct.models.model.MLModelAsset.from_memory(
+        spec_data=spec_data, blob_mapping={"weights/weight.bin": weights_data}
+    )
+    compiled_model = ct.models.CompiledMLModel.from_asset(asset=asset)
+    result = compiled_model.predict(
+      {
+          "x": np.array([1.0]),
+          "y": np.array([2.0]),
+      }
+    )    
+```

docs/conf.py

@@ -21,8 +21,8 @@
 project = 'coremltools API Reference'
 copyright = '2021, Apple Inc'
 author = 'Apple Inc'
-release = '8.0b1'
-version = '8.0b1'
+release = '8.1'
+version = '8.1'
 
 # -- General configuration ---------------------------------------------------
 

docs/source/coremltools.models.rst

@@ -80,3 +80,15 @@ utils
 
 .. automodule:: coremltools.models.utils
    :members:
+
+compute\_plan
+-------------------------------
+
+.. automodule:: coremltools.models.compute_plan
+   :members:
+
+compute\_device
+-------------------------------
+
+.. automodule:: coremltools.models.compute_device
+   :members: