From ccae25849587dda3dbdb4e6cbc8836cba506220e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Thu, 18 Apr 2024 09:48:26 +0200 Subject: [PATCH] Document `AppleOps` and `MPSOps` (#929) * Document AppleOps and MPSOps * Reformat Ops table - Sort alphabetically. - Note that `AppleOps` is new in 9.0. * Missing comma --- website/docs/api-backends.md | 22 ++++++++-------- website/docs/api-model.md | 50 ++++++++++++++++++------------------ 2 files changed, 37 insertions(+), 35 deletions(-) diff --git a/website/docs/api-backends.md b/website/docs/api-backends.md index fc69a775d..853fada3b 100644 --- a/website/docs/api-backends.md +++ b/website/docs/api-backends.md @@ -17,16 +17,18 @@ specialized versions can be called for different backends. You can also create your own `Ops` subclasses with specialized routines for your layers, and use the [`set_current_ops`](#set_current_ops) function to change the default. -| Backend | CPU | GPU | TPU | Description | -| ---------- | :----------------: | :----------------: | :---------------: | ----------------------------------------------------------------------------------------------------- | -| `NumpyOps` | | | | Execute via `numpy`, [`blis`](https://github.com/explosion/cython-blis) (optional) and custom Cython. | -| `CupyOps` | | | | Execute via [`cupy`](https://cupy.chainer.org/) and custom CUDA. | +| Backend | CPU | GPU | TPU | Description | +| ---------- | :----------------: | :----------------: | :---------------: | ----------------------------------------------------------------------------------------------------------- | +| `AppleOps` | | | | Use AMX matrix multiplication units on Apple Silicon Macs. Added in Thinc 9.0. | +| `CupyOps` | | | | Execute via [`cupy`](https://cupy.chainer.org/) and custom CUDA. | +| `MPSOps` | | | | Use the GPU on Apple Silicon Macs for PyTorch models, use AMX matrix multiplication units for Thinc Models. | +| `NumpyOps` | | | | Execute via `numpy`, [`blis`](https://github.com/explosion/cython-blis) (optional) and custom Cython. | ## Ops {#ops tag="class"} -The `Ops` class is typically not used directly but via `NumpyOps` or `CupyOps`, -which are subclasses of `Ops` and implement a **more efficient subset of the -methods**. You also have access to the ops via the +The `Ops` class is typically not used directly but via `NumpyOps`, `AppleOps`, +`CupyOps` or `MPSOps`, which are subclasses of `Ops` and implement a **more +efficient subset of the methods**. You also have access to the ops via the [`Model.ops`](/docs/api-model#attributes) attribute. The documented methods below list which backends provide optimized and more efficient versions (indicated by ), and which use the default implementation. @@ -56,7 +58,7 @@ use_ops(blis_ops) | Name | Type | Description | | ------------- | ------------ | ---------------------------------------------------------------------------------------- | -| `name` | str | **Class attribute:** Backend name, `"numpy"` or `"cupy"`. | +| `name` | str | **Class attribute:** Backend name, `"numpy"`, `"apple"`, `"cupy"` or `"mps"`. | | `xp` | Xp | **Class attribute:** `numpy` or `cupy`. | | `device_type` | str | The device type to use, if available for the given backend: `"cpu"`, `"gpu"` or `"tpu"`. | | `device_id` | int | The device ID to use, if available for the given backend. | @@ -1553,7 +1555,7 @@ numpy_ops = get_ops("numpy") | Argument | Type | Description | | ----------- | ------------ | ----------------------------------------------------- | -| `ops` | str | `"numpy"` or `"cupy"`. | +| `ops` | str | `"numpy"`, `"apple"`, `"cupy"` or `"mps"`. | | `**kwargs` | | Optional arguments passed to [`Ops.__init__`](#init). | | **RETURNS** | Ops | The backend object. | @@ -1572,7 +1574,7 @@ with use_ops("cupy"): | Argument | Type | Description | | ---------- | ------------ | ----------------------------------------------------- | -| `ops` | str | `"numpy"` or `"cupy"`. | +| `ops` | str | `"numpy"`, `"apple"`, `"cupy"` or `"mps"`. | | `**kwargs` | | Optional arguments passed to [`Ops.__init__`](#init). | ### get_current_ops {#get_current_ops tag="function"} diff --git a/website/docs/api-model.md b/website/docs/api-model.md index 597f67ec9..193fd1acb 100644 --- a/website/docs/api-model.md +++ b/website/docs/api-model.md @@ -84,19 +84,19 @@ model = Model( ) ``` -| Argument | Type | Description | -| -------------- | ------------------------------------------- | --------------------------------------------------------------------------------------- | -| `name` | str | The name of the layer type. | -| `forward` | Callable | Function to compute the forward result and the backpropagation callback. | -| _keyword-only_ | | | -| `init` | Callable | Function to define the initialization logic. | -| `dims` | Dict[str, Optional[int]] | Dictionary describing the model's dimensions. Map unknown dimensions to `None`. | -| `params` | Dict[str, Optional[FloatsXd]] | Dictionary with the model's parameters. Set currently unavailable parameters to `None`. | -| `refs` | Dict[str, Optional[Model]] | Dictionary mapping specific nodes (sublayers) of the network to a name. | -| `attrs` | Dict[str, Any] | Dictionary of non-parameter attributes. | -| `layers` | List[Model] | List of child layers. | -| `shims` | List[Shim] | List of interfaces for external models. | -| `ops` | Optional[Union[NumpyOps, CupyOps]] | An `Ops` instance, which provides mathematical and memory operations. | +| Argument | Type | Description | +| -------------- | ------------------------------------------------------------- | --------------------------------------------------------------------------------------- | +| `name` | str | The name of the layer type. | +| `forward` | Callable | Function to compute the forward result and the backpropagation callback. | +| _keyword-only_ | | | +| `init` | Callable | Function to define the initialization logic. | +| `dims` | Dict[str, Optional[int]] | Dictionary describing the model's dimensions. Map unknown dimensions to `None`. | +| `params` | Dict[str, Optional[FloatsXd]] | Dictionary with the model's parameters. Set currently unavailable parameters to `None`. | +| `refs` | Dict[str, Optional[Model]] | Dictionary mapping specific nodes (sublayers) of the network to a name. | +| `attrs` | Dict[str, Any] | Dictionary of non-parameter attributes. | +| `layers` | List[Model] | List of child layers. | +| `shims` | List[Shim] | List of interfaces for external models. | +| `ops` | Optional[Union[NumpyOps, AppleOps, CupyOps, MPSOps]] | An `Ops` instance, which provides mathematical and memory operations. | ### Model.define_operators {#define_operators tag="classmethod,contextmanager"} @@ -260,17 +260,17 @@ for node in model.walk(): The `walk` method supports three iteration orders through the `order` argument: -* `"bfs"`: breadth-first. Iteration order of the example above: - *1 - 2 - 4 - 3 - 5* -* `"dfs_pre"`: depth-first preorder, outputs a node before its children. - Iteration order of the example above: *1 - 2 - 3 - 4 - 5* -* `"dfs_post"`: depth-first postorder, outputs children before a node itself. - Iteration order of the example above: *3 - 2 - 5 - 4 - 1* +- `"bfs"`: breadth-first. Iteration order of the example above: _1 - 2 - 4 - 3 - + 5_ +- `"dfs_pre"`: depth-first preorder, outputs a node before its children. + Iteration order of the example above: _1 - 2 - 3 - 4 - 5_ +- `"dfs_post"`: depth-first postorder, outputs children before a node itself. + Iteration order of the example above: _3 - 2 - 5 - 4 - 1_ -| Argument | Type | Description | -|-------------|--------------------------|--------------------------------------------------------------------------------------------------------------------------------------------| +| Argument | Type | Description | +| ----------- | ------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------- | | `order` | str | Node iteration order. `"bfs"` (breadth-first), `"dfs_pre"` (depth-first preorder), `"dfs_post"` (depth-first postorder) Default: `"bfs"`. | -| **RETURNS** | Iterable[Model] | The layers of the model. | +| **RETURNS** | Iterable[Model] | The layers of the model. | ### Model.remove_node {#remove_node tag="method"} @@ -329,9 +329,9 @@ assert model.get_dim("nI") == 16 Retrieve the value of a dimension of the given name, or `None` if the dimension is either unregistered or the value is currently unset. -| Argument | Type | Description | -| ----------- | --------------------- | --------------------------------------- | -| `name` | str | The name of the dimension, e.g. `"nO"`. | +| Argument | Type | Description | +| ----------- | ---------------------- | --------------------------------------- | +| `name` | str | The name of the dimension, e.g. `"nO"`. | | **RETURNS** | Optional[int] | The size of the dimension, or `None`. | ### Model.set_dim {#set_dim tag="method"}