[dnn] Improve error handling in dnn/linear module (#704)

Improve error handling in `dnn/linear` module

Author: Routhleck

brainpy/_src/dnn/linear.py

@@ -12,7 +12,7 @@
 from brainpy import math as bm
 from brainpy._src import connect, initialize as init
 from brainpy._src.context import share
-from brainpy._src.dependency_check import import_taichi, import_braintaichi
+from brainpy._src.dependency_check import import_taichi, import_braintaichi, raise_braintaichi_not_found
 from brainpy._src.dnn.base import Layer
 from brainpy._src.mixin import SupportOnline, SupportOffline, SupportSTDP
 from brainpy.check import is_initializer
@@ -241,7 +241,7 @@ def update(self, x):
         return x
 
 
-if ti is not None:
+if ti is not None and bti is not None:
 
     # @numba.njit(nogil=True, fastmath=True, parallel=False)
     # def _cpu_dense_on_post(weight, spike, trace, w_min, w_max, out_w):
@@ -321,7 +321,7 @@ def _dense_on_pre(
 
 def dense_on_pre(weight, spike, trace, w_min, w_max):
     if dense_on_pre_prim is None:
-        raise PackageMissingError.by_purpose('taichi', 'custom operators')
+        raise_braintaichi_not_found()
 
     if w_min is None:
         w_min = -np.inf
@@ -341,7 +341,7 @@ def dense_on_pre(weight, spike, trace, w_min, w_max):
 
 def dense_on_post(weight, spike, trace, w_min, w_max):
     if dense_on_post_prim is None:
-        raise PackageMissingError.by_purpose('taichi', 'custom operators')
+        raise_braintaichi_not_found()
 
     if w_min is None:
         w_min = -np.inf
@@ -728,7 +728,7 @@ def _batch_csrmv(self, x):
                               transpose=self.transpose)
 
 
-if ti is not None:
+if ti is not None and bti is not None:
     @ti.kernel
     def _csr_on_pre_update(
         old_w: ti.types.ndarray(ndim=1),  # vector with shape of (num_syn)
@@ -852,7 +852,7 @@ def _csc_on_post_update(
 
 def csr_on_pre_update(w, indices, indptr, spike, trace, w_min=None, w_max=None):
     if csr_on_pre_update_prim is None:
-        raise PackageMissingError.by_purpose('taichi', 'customized operators')
+        raise_braintaichi_not_found()
 
     if w_min is None:
         w_min = -np.inf
@@ -874,7 +874,7 @@ def csr_on_pre_update(w, indices, indptr, spike, trace, w_min=None, w_max=None):
 
 def coo_on_pre_update(w, pre_ids, post_ids, spike, trace, w_min=None, w_max=None):
     if coo_on_pre_update_prim is None:
-        raise PackageMissingError.by_purpose('taichi', 'customized operators')
+        raise_braintaichi_not_found()
 
     if w_min is None:
         w_min = -np.inf
@@ -897,7 +897,7 @@ def coo_on_pre_update(w, pre_ids, post_ids, spike, trace, w_min=None, w_max=None
 
 def csc_on_post_update(w, post_ids, indptr, w_ids, post_spike, pre_trace, w_min=None, w_max=None):
     if csc_on_post_update_prim is None:
-        raise PackageMissingError.by_purpose('taichi', 'customized operators')
+        raise_braintaichi_not_found()
 
     if w_min is None:
         w_min = -np.inf

docker/requirements.txt

@@ -6,10 +6,10 @@ jax
 jaxlib
 scipy>=1.1.0
 brainpy
-brainpylib
 brainpy_datasets
 h5py
 pathos
+braintaichi
 
 # test requirements
 pytest

docs/quickstart/installation.rst

@@ -36,7 +36,6 @@ To install brainpy with minimum requirements (only depends on ``jax``), you can
 
     # or
 
-    pip install brainpy[cuda11_mini] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  # for CUDA 11.0
     pip install brainpy[cuda12_mini] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  # for CUDA 12.0
 
     # or
@@ -64,7 +63,6 @@ To install a GPU-only version of BrainPy, you can run
 .. code-block:: bash
 
     pip install brainpy[cuda12] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  # for CUDA 12.0
-    pip install brainpy[cuda11] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  # for CUDA 11.0
 
 
 
@@ -79,25 +77,3 @@ you can run the following in your cloud TPU VM:
     pip install brainpy[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html  # for google TPU
 
 
-
-``brainpylib``
---------------
-
-
-``brainpylib`` defines a set of useful operators for building and simulating spiking neural networks.
-
-
-To install the ``brainpylib`` package on CPU devices, you can run
-
-.. code-block:: bash
-
-    pip install brainpylib
-
-
-To install the ``brainpylib`` package on CUDA (Linux only), you can run
-
-
-.. code-block:: bash
-
-    pip install brainpylib
-