Skip to content

Commit

Permalink
Document AppleOps and MPSOps (#929)
Browse files Browse the repository at this point in the history 
* Document AppleOps and MPSOps

* Reformat Ops table

- Sort alphabetically.
- Note that `AppleOps` is new in 9.0.

* Missing comma
  • Loading branch information
@danieldk
danieldk committed Apr 18, 2024
1 parent 2a0b9c1 commit ccae258
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 35 deletions.
22 changes: 12 additions & 10 deletions website/docs/api-backends.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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` | <i name="yes"></i> | <i name="no"></i> | <i name="no"></i> | Execute via `numpy`, [`blis`](https://github.com/explosion/cython-blis) (optional) and custom Cython. |
| `CupyOps` | <i name="no"></i> | <i name="yes"></i> | <i name="no"></i> | Execute via [`cupy`](https://cupy.chainer.org/) and custom CUDA. |
| Backend | CPU | GPU | TPU | Description |
| ---------- | :----------------: | :----------------: | :---------------: | ----------------------------------------------------------------------------------------------------------- |
| `AppleOps` | <i name="yes"></i> | <i name="no"></i> | <i name="no"></i> | Use AMX matrix multiplication units on Apple Silicon Macs. Added in Thinc 9.0. |
| `CupyOps` | <i name="no"></i> | <i name="yes"></i> | <i name="no"></i> | Execute via [`cupy`](https://cupy.chainer.org/) and custom CUDA. |
| `MPSOps` | <i name="yes"></i> | <i name="yes"></i> | <i name="no"></i> | Use the GPU on Apple Silicon Macs for PyTorch models, use AMX matrix multiplication units for Thinc Models. |
| `NumpyOps` | <i name="yes"></i> | <i name="no"></i> | <i name="no"></i> | 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 <i name="yes"></i>), and which use the default implementation.
Expand Down Expand Up @@ -56,7 +58,7 @@ use_ops(blis_ops)

| Name | Type | Description |
| ------------- | ------------ | ---------------------------------------------------------------------------------------- |
| `name` | <tt>str</tt> | **Class attribute:** Backend name, `"numpy"` or `"cupy"`. |
| `name` | <tt>str</tt> | **Class attribute:** Backend name, `"numpy"`, `"apple"`, `"cupy"` or `"mps"`. |
| `xp` | <tt>Xp</tt> | **Class attribute:** `numpy` or `cupy`. |
| `device_type` | <tt>str</tt> | The device type to use, if available for the given backend: `"cpu"`, `"gpu"` or `"tpu"`. |
| `device_id` | <tt>int</tt> | The device ID to use, if available for the given backend. |
Expand Down Expand Up @@ -1553,7 +1555,7 @@ numpy_ops = get_ops("numpy")

| Argument | Type | Description |
| ----------- | ------------ | ----------------------------------------------------- |
| `ops` | <tt>str</tt> | `"numpy"` or `"cupy"`. |
| `ops` | <tt>str</tt> | `"numpy"`, `"apple"`, `"cupy"` or `"mps"`. |
| `**kwargs` | | Optional arguments passed to [`Ops.__init__`](#init). |
| **RETURNS** | <tt>Ops</tt> | The backend object. |

Expand All @@ -1572,7 +1574,7 @@ with use_ops("cupy"):

| Argument | Type | Description |
| ---------- | ------------ | ----------------------------------------------------- |
| `ops` | <tt>str</tt> | `"numpy"` or `"cupy"`. |
| `ops` | <tt>str</tt> | `"numpy"`, `"apple"`, `"cupy"` or `"mps"`. |
| `**kwargs` | | Optional arguments passed to [`Ops.__init__`](#init). |

### get_current_ops {#get_current_ops tag="function"}
Expand Down
50 changes: 25 additions & 25 deletions website/docs/api-model.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,19 +84,19 @@ model = Model(
)
```

| Argument | Type | Description |
| -------------- | ------------------------------------------- | --------------------------------------------------------------------------------------- |
| `name` | <tt>str</tt> | The name of the layer type. |
| `forward` | <tt>Callable</tt> | Function to compute the forward result and the backpropagation callback. |
| _keyword-only_ | | |
| `init` | <tt>Callable</tt> | Function to define the initialization logic. |
| `dims` | <tt>Dict[str, Optional[int]]</tt> | Dictionary describing the model's dimensions. Map unknown dimensions to `None`. |
| `params` | <tt>Dict[str, Optional[FloatsXd]]</tt> | Dictionary with the model's parameters. Set currently unavailable parameters to `None`. |
| `refs` | <tt>Dict[str, Optional[Model]]</tt> | Dictionary mapping specific nodes (sublayers) of the network to a name. |
| `attrs` | <tt>Dict[str, Any]</tt> | Dictionary of non-parameter attributes. |
| `layers` | <tt>List[Model]</tt> | List of child layers. |
| `shims` | <tt>List[Shim]</tt> | List of interfaces for external models. |
| `ops` | <tt>Optional[Union[NumpyOps, CupyOps]]</tt> | An `Ops` instance, which provides mathematical and memory operations. |
| Argument | Type | Description |
| -------------- | ------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
| `name` | <tt>str</tt> | The name of the layer type. |
| `forward` | <tt>Callable</tt> | Function to compute the forward result and the backpropagation callback. |
| _keyword-only_ | | |
| `init` | <tt>Callable</tt> | Function to define the initialization logic. |
| `dims` | <tt>Dict[str, Optional[int]]</tt> | Dictionary describing the model's dimensions. Map unknown dimensions to `None`. |
| `params` | <tt>Dict[str, Optional[FloatsXd]]</tt> | Dictionary with the model's parameters. Set currently unavailable parameters to `None`. |
| `refs` | <tt>Dict[str, Optional[Model]]</tt> | Dictionary mapping specific nodes (sublayers) of the network to a name. |
| `attrs` | <tt>Dict[str, Any]</tt> | Dictionary of non-parameter attributes. |
| `layers` | <tt>List[Model]</tt> | List of child layers. |
| `shims` | <tt>List[Shim]</tt> | List of interfaces for external models. |
| `ops` | <tt>Optional[Union[NumpyOps, AppleOps, CupyOps, MPSOps]]</tt> | An `Ops` instance, which provides mathematical and memory operations. |

### Model.define_operators {#define_operators tag="classmethod,contextmanager"}

Expand Down Expand Up @@ -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` | <tt>str</tt> | Node iteration order. `"bfs"` (breadth-first), `"dfs_pre"` (depth-first preorder), `"dfs_post"` (depth-first postorder) Default: `"bfs"`. |
| **RETURNS** | <tt>Iterable[Model]</tt> | The layers of the model. |
| **RETURNS** | <tt>Iterable[Model]</tt> | The layers of the model. |

### Model.remove_node {#remove_node tag="method"}

Expand Down Expand Up @@ -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` | <tt>str</tt> | The name of the dimension, e.g. `"nO"`. |
| Argument | Type | Description |
| ----------- | ---------------------- | --------------------------------------- |
| `name` | <tt>str</tt> | The name of the dimension, e.g. `"nO"`. |
| **RETURNS** | <tt>Optional[int]</tt> | The size of the dimension, or `None`. |

### Model.set_dim {#set_dim tag="method"}
Expand Down

0 comments on commit ccae258

Please sign in to comment.