compatible with jax>=0.8.0

Author: chaoming0625

brainpy/__init__.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "2.7.1"
-__version_info__ = (2, 7, 1)
+__version__ = "2.7.2"
+__version_info__ = (2, 7, 2)
 
 
 from brainpy import _errors as errors

brainpy/math/object_transform/base.py

@@ -30,14 +30,10 @@
 from brainpy.math.defaults import defaults
 from brainpy.math.modes import Mode
 from brainpy.math.ndarray import (Array, )
-from brainpy.math.object_transform.collectors import (ArrayCollector, Collector)
-from brainpy.math.object_transform.naming import (
-    get_unique_name,
-    check_name_uniqueness
-)
+from brainpy.math.object_transform.collectors import ArrayCollector, Collector
+from brainpy.math.object_transform.naming import get_unique_name, check_name_uniqueness
 from brainpy.math.object_transform.variables import (
-    Variable, VariableView, TrainVar,
-    VarList, VarDict
+    Variable, VariableView, TrainVar, VarList, VarDict
 )
 from brainpy.math.sharding import BATCH_AXIS
 

brainpy/math/object_transform/variables.py

@@ -14,17 +14,17 @@
 # ==============================================================================
 from typing import Optional, Any, Sequence
 
+import brainstate
 import jax
 import numpy as np
+from brainstate._state import record_state_value_read, record_state_value_write
 from jax import numpy as jnp
 from jax.dtypes import canonicalize_dtype
 from jax.tree_util import register_pytree_node_class
 
-import brainstate
 from brainpy._errors import MathError
 from brainpy.math.ndarray import Array
 from brainpy.math.sharding import BATCH_AXIS
-from brainstate._state import record_state_value_read, record_state_value_write
 
 __all__ = [
     'Variable',

brainpy/math/sparse/utils.py

@@ -13,57 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import warnings
-from functools import partial
+
 from typing import Tuple
 
-import jax
-import numpy as np
-from jax import core, numpy as jnp
-from jax import lax
-from jax.interpreters import batching
-from jax.interpreters import mlir, ad
-from jax.tree_util import tree_flatten, tree_unflatten
-from jaxlib import gpu_sparse
+from jax import numpy as jnp
+from jax.experimental.sparse import csr_todense
 
 from brainpy.math.interoperability import as_jax
 
-if jax.__version__ >= '0.5.0':
-    from jax.extend.core import Primitive
-else:
-    from jax.core import Primitive
-
 __all__ = [
     'coo_to_csr',
     'csr_to_coo',
     'csr_to_dense'
 ]
 
 
-def _general_batching_rule(prim, args, axes, **kwargs):
-    batch_axes, batch_args, non_batch_args = [], {}, {}
-    for ax_i, ax in enumerate(axes):
-        if ax is None:
-            non_batch_args[f'ax{ax_i}'] = args[ax_i]
-        else:
-            batch_args[f'ax{ax_i}'] = args[ax_i] if ax == 0 else jnp.moveaxis(args[ax_i], ax, 0)
-            batch_axes.append(ax_i)
-
-    def f(_, x):
-        pars = tuple([(x[f'ax{i}'] if i in batch_axes else non_batch_args[f'ax{i}'])
-                      for i in range(len(axes))])
-        return 0, prim.bind(*pars, **kwargs)
-
-    _, outs = lax.scan(f, 0, batch_args)
-    out_vals, out_tree = tree_flatten(outs)
-    out_dim = tree_unflatten(out_tree, (0,) * len(out_vals))
-    return outs, out_dim
-
-
-def _register_general_batching(prim):
-    batching.primitive_batchers[prim] = partial(_general_batching_rule, prim)
-
-
 def coo_to_csr(
     pre_ids: jnp.ndarray,
     post_ids: jnp.ndarray,
@@ -97,108 +61,6 @@ def csr_to_coo(
     return jnp.cumsum(jnp.zeros_like(indices).at[indptr].add(1)) - 1, indices
 
 
-def csr_to_csc():
-    pass
-
-
-def coo_to_dense(
-    data: jnp.ndarray,
-    rows: jnp.ndarray,
-    cols: jnp.ndarray,
-    *,
-    shape: Tuple[int, int]
-) -> jnp.ndarray:
-    pass
-
-
-def csr_to_dense(
-    data: jnp.ndarray,
-    indices: jnp.ndarray,
-    indptr: jnp.ndarray,
-    *,
-    shape: Tuple[int, int]
-) -> jnp.ndarray:
-    data = as_jax(data)
-    indices = as_jax(indices)
-    indptr = as_jax(indptr)
-    return csr_to_dense_p.bind(data, indices, indptr, shape=shape)
-
-
-def _coo_extract(row, col, mat):
-    """Extract values of dense matrix mat at given COO indices."""
-    return mat[row, col]
-
-
-def _csr_extract(indices, indptr, mat):
-    """Extract values of dense matrix mat at given CSR indices."""
-    return _coo_extract(*csr_to_coo(indices, indptr), mat)
-
-
-def _coo_todense(data, row, col, *, shape):
-    """Convert CSR-format sparse matrix to a dense matrix.
-
-    Args:
-      data : array of shape ``(nse,)``.
-      row : array of shape ``(nse,)``
-      col : array of shape ``(nse,)`` and dtype ``row.dtype``
-      shape : COOInfo object containing matrix metadata
-
-    Returns:
-      mat : array with specified shape and dtype matching ``data``
-    """
-    return jnp.zeros(shape, data.dtype).at[row, col].add(data)
-
-
-def _csr_to_dense_impl(data, indices, indptr, *, shape):
-    return _coo_todense(data, *csr_to_coo(indices, indptr), shape=shape)
-
-
-def _csr_to_dense_abstract_eval(data, indices, indptr, *, shape):
-    assert data.ndim == indices.ndim == indptr.ndim == 1
-    assert indices.dtype == indptr.dtype
-    assert data.shape == indices.shape
-    assert indptr.shape[0] == shape[0] + 1
-    return core.ShapedArray(shape, data.dtype)
-
-
-_csr_to_dense_lowering = mlir.lower_fun(_csr_to_dense_impl, multiple_results=False)
-
-
-def _csr_to_dense_gpu_lowering(ctx, data, indices, indptr, *, shape):
-    data_aval, indices_aval, _ = ctx.avals_in
-    dtype = data_aval.dtype
-    if not (np.issubdtype(dtype, np.floating) or np.issubdtype(dtype, np.complexfloating)):
-        warnings.warn(f"csr_todense cusparse/hipsparse lowering not available for dtype={dtype}. "
-                      "Falling back to default implementation.",
-                      UserWarning)
-        return _csr_to_dense_lowering(ctx, data, indices, indptr, shape=shape)
-    return [gpu_sparse.cuda_csr_todense(data, indices, indptr,
-                                        shape=shape, data_dtype=dtype,
-                                        index_dtype=indices_aval.dtype)]
-
-
-def _csr_to_dense_jvp(data_dot, data, indices, indptr, *, shape):
-    return csr_to_dense(data_dot, indices, indptr, shape=shape)
-
-
-def _csr_to_dense_transpose(ct, data, indices, indptr, *, shape):
-    # Note: we assume that transpose has the same sparsity pattern.
-    # Can we check this?
-    assert ad.is_undefined_primal(data)
-    if ad.is_undefined_primal(indices) or ad.is_undefined_primal(indptr):
-        raise ValueError("Cannot transpose with respect to sparse indices")
-    assert ct.shape == shape
-    assert indices.aval.dtype == indptr.aval.dtype
-    assert ct.dtype == data.aval.dtype
-    return _csr_extract(indices, indptr, ct), indices, indptr
+csr_to_dense = csr_todense
 
 
-csr_to_dense_p = Primitive('csr_to_dense')
-csr_to_dense_p.def_impl(_csr_to_dense_impl)
-csr_to_dense_p.def_abstract_eval(_csr_to_dense_abstract_eval)
-ad.defjvp(csr_to_dense_p, _csr_to_dense_jvp, None, None)
-ad.primitive_transposes[csr_to_dense_p] = _csr_to_dense_transpose
-mlir.register_lowering(csr_to_dense_p, _csr_to_dense_lowering)
-_register_general_batching(csr_to_dense_p)
-if gpu_sparse.cuda_is_supported:
-    mlir.register_lowering(csr_to_dense_p, _csr_to_dense_gpu_lowering, platform='cuda')

docs_state/index.rst

@@ -117,7 +117,7 @@ Learn more
    .. grid-item::
       :columns: 6 6 6 4
 
-      .. card:: :material-regular:`data_exploration;2em` Classical APIs
+      .. card:: :material-regular:`data_exploration;2em` ``brainpy`` APIs
          :class-card: sd-text-black sd-bg-light
          :link: https://brainpy.readthedocs.io/