Import mxnet and tensorflow only if explicitly enabled (#890)

* Import mxnet and tensorflow only if explicitly enabled

* Ignore import errors for mxnet/tensorflow in tests

* Add enable_{mxnet,tensorflow} to thinc.api and docs

* Update intro example notebook

* Add warnings/info to docs

* Add deprecation warnings to enable_ methods

* Extend error messages in assert_{mxnet,tensorflow}_installed

Author: adrianeboyd

.github/workflows/tests.yml

@@ -136,7 +136,7 @@ jobs:
         shell: bash --noprofile --norc -o pipefail {0}
 
       - name: Run tests with extras
-        run: python -m pytest --pyargs thinc --cov=thinc --cov-report=term
+        run: python -m pytest --pyargs thinc --cov=thinc --cov-report=term -p thinc.tests.enable_tensorflow -p thinc.tests.enable_mxnet
 
       - name: Run tests for thinc-apple-ops
         run: |

examples/00_intro_to_thinc.ipynb

@@ -23,7 +23,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip install \"thinc>=8.0.0\" \"ml_datasets>=0.2.0\" \"tqdm>=4.41\""
+    "!pip install \"thinc>=8.2.0\" \"ml_datasets>=0.2.0\" \"tqdm>=4.41\""
    ]
   },
   {
@@ -1050,7 +1050,8 @@
    "source": [
     "from tensorflow.keras.layers import Dense, Dropout\n",
     "from tensorflow.keras.models import Sequential\n",
-    "from thinc.api import TensorFlowWrapper, Adam\n",
+    "from thinc.api import enable_tensorflow, TensorFlowWrapper, Adam\n",
+    "enable_tensorflow()\n",
     "\n",
     "width = 32\n",
     "nO = 10\n",
@@ -1373,8 +1374,9 @@
    "outputs": [],
    "source": [
     "from mxnet.gluon.nn import Dense, Sequential, Dropout\n",
-    "from thinc.api import MXNetWrapper, chain, Softmax\n",
+    "from thinc.api import enable_mxnet, MXNetWrapper, chain, Softmax\n",
     "import thinc.util\n",
+    "enable_mxnet()\n",
     "\n",
     "assert thinc.util.has_mxnet\n",
     "\n",

thinc/api.py

@@ -11,7 +11,7 @@
     use_pytorch_for_gpu_memory,
     use_tensorflow_for_gpu_memory,
 )
-from .compat import has_cupy
+from .compat import enable_mxnet, enable_tensorflow, has_cupy
 from .config import Config, ConfigValidationError, registry
 from .initializers import (
     configure_normal_init,
@@ -190,6 +190,8 @@
     "torch2xp", "xp2torch", "tensorflow2xp", "xp2tensorflow", "mxnet2xp", "xp2mxnet",
     "get_torch_default_device",
     # .compat
+    "enable_mxnet",
+    "enable_tensorflow",
     "has_cupy",
     # .backends
     "get_ops", "set_current_ops", "get_current_ops", "use_ops",

thinc/backends/_cupy_allocators.py

@@ -12,7 +12,7 @@ def cupy_tensorflow_allocator(size_in_bytes: int):
     sitting in the other library's pool.
     """
     size_in_bytes = max(1024, size_in_bytes)
-    tensor = tensorflow.zeros((size_in_bytes // 4,), dtype=tensorflow.dtypes.float32)
+    tensor = tensorflow.zeros((size_in_bytes // 4,), dtype=tensorflow.dtypes.float32)  # type: ignore
     # We convert to cupy via dlpack, so that we can get a memory pointer.
     cupy_array = cast(ArrayXd, tensorflow2xp(tensor))
     address = int(cupy_array.data)

thinc/compat.py

@@ -1,3 +1,5 @@
+import warnings
+
 from packaging.version import Version
 
 try:  # pragma: no cover
@@ -50,25 +52,45 @@
     has_torch_amp = False
     torch_version = Version("0.0.0")
 
-try:  # pragma: no cover
+
+def enable_tensorflow():
+    warn_msg = (
+        "Built-in TensorFlow support will be removed in Thinc v9. If you need "
+        "TensorFlow support in the future, you can transition to using a "
+        "custom copy of the current TensorFlowWrapper in your package or "
+        "project."
+    )
+    warnings.warn(warn_msg, DeprecationWarning)
+    global tensorflow, has_tensorflow, has_tensorflow_gpu
     import tensorflow
     import tensorflow.experimental.dlpack
 
     has_tensorflow = True
     has_tensorflow_gpu = len(tensorflow.config.get_visible_devices("GPU")) > 0
-except ImportError:  # pragma: no cover
-    tensorflow = None
-    has_tensorflow = False
-    has_tensorflow_gpu = False
 
 
-try:  # pragma: no cover
+tensorflow = None
+has_tensorflow = False
+has_tensorflow_gpu = False
+
+
+def enable_mxnet():
+    warn_msg = (
+        "Built-in MXNet support will be removed in Thinc v9. If you need "
+        "MXNet support in the future, you can transition to using a "
+        "custom copy of the current MXNetWrapper in your package or "
+        "project."
+    )
+    warnings.warn(warn_msg, DeprecationWarning)
+    global mxnet, has_mxnet
     import mxnet
 
     has_mxnet = True
-except ImportError:  # pragma: no cover
-    mxnet = None
-    has_mxnet = False
+
+
+mxnet = None
+has_mxnet = False
+
 
 try:
     import h5py

thinc/layers/tensorflowwrapper.py

@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 from typing import Any, Callable, Dict, Optional, Tuple, Type, TypeVar
 
 import srsly

thinc/shims/mxnet.py

@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 import copy
 from typing import Any, cast
 

thinc/shims/tensorflow.py

@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 import contextlib
 import copy
 from io import BytesIO

thinc/tests/enable_mxnet.py

@@ -0,0 +1,6 @@
+from thinc.compat import enable_mxnet
+
+try:
+    enable_mxnet()
+except ImportError:
+    pass

thinc/tests/enable_tensorflow.py

@@ -0,0 +1,6 @@
+from thinc.compat import enable_tensorflow
+
+try:
+    enable_tensorflow()
+except ImportError:
+    pass

thinc/util.py

@@ -151,7 +151,7 @@ def is_torch_mps_array(obj: Any) -> bool:  # pragma: no cover
 def is_tensorflow_array(obj: Any) -> bool:  # pragma: no cover
     if not has_tensorflow:
         return False
-    elif isinstance(obj, tf.Tensor):
+    elif isinstance(obj, tf.Tensor):  # type: ignore
         return True
     else:
         return False
@@ -164,7 +164,7 @@ def is_tensorflow_gpu_array(obj: Any) -> bool:  # pragma: no cover
 def is_mxnet_array(obj: Any) -> bool:  # pragma: no cover
     if not has_mxnet:
         return False
-    elif isinstance(obj, mx.nd.NDArray):
+    elif isinstance(obj, mx.nd.NDArray):  # type: ignore
         return True
     else:
         return False
@@ -316,15 +316,17 @@ def get_width(
 
 def assert_tensorflow_installed() -> None:  # pragma: no cover
     """Raise an ImportError if TensorFlow is not installed."""
-    template = "TensorFlow support requires {pkg}: pip install thinc[tensorflow]"
+    template = "TensorFlow support requires {pkg}: pip install thinc[tensorflow]\n\nEnable TensorFlow support with thinc.api.enable_tensorflow()"
     if not has_tensorflow:
-        raise ImportError(template.format(pkg="tensorflow>=2.0.0"))
+        raise ImportError(template.format(pkg="tensorflow>=2.0.0,<2.6.0"))
 
 
 def assert_mxnet_installed() -> None:  # pragma: no cover
     """Raise an ImportError if MXNet is not installed."""
     if not has_mxnet:
-        raise ImportError("MXNet support requires mxnet: pip install thinc[mxnet]")
+        raise ImportError(
+            "MXNet support requires mxnet: pip install thinc[mxnet]\n\nEnable MXNet support with thinc.api.enable_mxnet()"
+        )
 
 
 def assert_pytorch_installed() -> None:  # pragma: no cover
@@ -429,32 +431,32 @@ def torch2xp(
 
 def xp2tensorflow(
     xp_tensor: ArrayXd, requires_grad: bool = False, as_variable: bool = False
-) -> "tf.Tensor":  # pragma: no cover
+) -> "tf.Tensor":  # type: ignore  # pragma: no cover
     """Convert a numpy or cupy tensor to a TensorFlow Tensor or Variable"""
     assert_tensorflow_installed()
     if hasattr(xp_tensor, "toDlpack"):
         dlpack_tensor = xp_tensor.toDlpack()  # type: ignore
-        tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_tensor)
+        tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_tensor)  # type: ignore
     elif hasattr(xp_tensor, "__dlpack__"):
         dlpack_tensor = xp_tensor.__dlpack__()  # type: ignore
-        tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_tensor)
+        tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_tensor)  # type: ignore
     else:
-        tf_tensor = tf.convert_to_tensor(xp_tensor)
+        tf_tensor = tf.convert_to_tensor(xp_tensor)  # type: ignore
     if as_variable:
         # tf.Variable() automatically puts in GPU if available.
         # So we need to control it using the context manager
-        with tf.device(tf_tensor.device):
-            tf_tensor = tf.Variable(tf_tensor, trainable=requires_grad)
+        with tf.device(tf_tensor.device):  # type: ignore
+            tf_tensor = tf.Variable(tf_tensor, trainable=requires_grad)  # type: ignore
     if requires_grad is False and as_variable is False:
         # tf.stop_gradient() automatically puts in GPU if available.
         # So we need to control it using the context manager
-        with tf.device(tf_tensor.device):
-            tf_tensor = tf.stop_gradient(tf_tensor)
+        with tf.device(tf_tensor.device):  # type: ignore
+            tf_tensor = tf.stop_gradient(tf_tensor)  # type: ignore
     return tf_tensor
 
 
 def tensorflow2xp(
-    tf_tensor: "tf.Tensor", *, ops: Optional["Ops"] = None
+    tf_tensor: "tf.Tensor", *, ops: Optional["Ops"] = None  # type: ignore
 ) -> ArrayXd:  # pragma: no cover
     """Convert a Tensorflow tensor to numpy or cupy tensor depending on the `ops` parameter.
     If `ops` is `None`, the type of the resultant tensor will be determined by the source tensor's device.
@@ -466,7 +468,7 @@ def tensorflow2xp(
         if isinstance(ops, NumpyOps):
             return tf_tensor.numpy()
         else:
-            dlpack_tensor = tf.experimental.dlpack.to_dlpack(tf_tensor)
+            dlpack_tensor = tf.experimental.dlpack.to_dlpack(tf_tensor)  # type: ignore
             return cupy_from_dlpack(dlpack_tensor)
     else:
         if isinstance(ops, NumpyOps) or ops is None:
@@ -477,21 +479,21 @@ def tensorflow2xp(
 
 def xp2mxnet(
     xp_tensor: ArrayXd, requires_grad: bool = False
-) -> "mx.nd.NDArray":  # pragma: no cover
+) -> "mx.nd.NDArray":  # type: ignore  # pragma: no cover
     """Convert a numpy or cupy tensor to a MXNet tensor."""
     assert_mxnet_installed()
     if hasattr(xp_tensor, "toDlpack"):
         dlpack_tensor = xp_tensor.toDlpack()  # type: ignore
-        mx_tensor = mx.nd.from_dlpack(dlpack_tensor)
+        mx_tensor = mx.nd.from_dlpack(dlpack_tensor)  # type: ignore
     else:
-        mx_tensor = mx.nd.from_numpy(xp_tensor)
+        mx_tensor = mx.nd.from_numpy(xp_tensor)  # type: ignore
     if requires_grad:
         mx_tensor.attach_grad()
     return mx_tensor
 
 
 def mxnet2xp(
-    mx_tensor: "mx.nd.NDArray", *, ops: Optional["Ops"] = None
+    mx_tensor: "mx.nd.NDArray", *, ops: Optional["Ops"] = None  # type: ignore
 ) -> ArrayXd:  # pragma: no cover
     """Convert a MXNet tensor to a numpy or cupy tensor."""
     from .api import NumpyOps

website/docs/api-layers.md

@@ -1003,7 +1003,7 @@ model, e.g. `chain(f, g)` computes `g(f(x))`.
 
 | Argument    | Type           | Description                       |
 | ----------- | -------------- | --------------------------------- |
-| `layer1 `   | <tt>Model</tt> | The first model to compose.       |
+| `layer1`    | <tt>Model</tt> | The first model to compose.       |
 | `layer2`    | <tt>Model</tt> | The second model to compose.      |
 | `*layers`   | <tt>Model</tt> | Any additional models to compose. |
 | **RETURNS** | <tt>Model</tt> | The composed feed-forward model.  |
@@ -1795,6 +1795,16 @@ https://github.com/explosion/thinc/blob/master/thinc/layers/torchscriptwrapper.p
 
 </inline-list>
 
+<infobox variant="warning">
+In Thinc v8.2+, TensorFlow support is not enabled by default. To enable TensorFlow:
+
+```python
+from thinc.api import enable_tensorflow
+enable_tensorflow()
+```
+
+</infobox>
+
 Wrap a [TensorFlow](https://tensorflow.org) model, so that it has the same API
 as Thinc models. To optimize the model, you'll need to create a TensorFlow
 optimizer and call `optimizer.apply_gradients` after each batch. To allow
@@ -1820,6 +1830,16 @@ https://github.com/explosion/thinc/blob/master/thinc/layers/tensorflowwrapper.py
 
 </inline-list>
 
+<infobox variant="warning">
+In Thinc v8.2+, MXNet support is not enabled by default. To enable MXNet:
+
+```python
+from thinc.api import enable_mxnet
+enable_mxnet()
+```
+
+</infobox>
+
 Wrap a [MXNet](https://mxnet.apache.org/) model, so that it has the same API as
 Thinc models. To optimize the model, you'll need to create a MXNet optimizer and
 call `optimizer.step()` after each batch. To allow maximum flexibility, the

website/docs/api-util.md

@@ -141,6 +141,14 @@ Converts a class vector (integers) to binary class matrix. Based on
 | `label_smoothing` | <tt>float</tt>         | Smoothing-coefficient for label-smoothing.                                                     |
 | **RETURNS**       | <tt>Floats2d</tt>      | A binary matrix representation of the input. The axis representing the classes is placed last. |
 
+### enable_mxnet {#enable_mxnet tag="function" new="8.2.0"}
+
+Import and enable internal support for MXNet.
+
+### enable_tensorflow {#enable_tensorflow tag="function"  new="8.2.0"}
+
+Import and enable internal support for TensorFlow.
+
 ### xp2torch {#xp2torch tag="function"}
 
 Convert a `numpy` or `cupy` tensor to a PyTorch tensor.

website/docs/usage-frameworks.md

@@ -81,6 +81,16 @@ Y, backprop = model(X, is_train=True)
 dX = backprop(Y)
 ```
 
+<infobox variant="warning">
+In Thinc v8.2+, TensorFlow support is not enabled by default. To enable TensorFlow:
+
+```python
+from thinc.api import enable_tensorflow
+enable_tensorflow()
+```
+
+</infobox>
+
 ```python
 ### TensorFlow Example {highlight="6"}
 from thinc.api import TensorFlowWrapper, chain, Linear