Source code for jax.experimental.sparse.random

# Copyright 2021 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import operator

from jax import dtypes
from jax import vmap
from jax import random
from jax.util import split_list
import jax.numpy as jnp
from jax.experimental import sparse


[docs] def random_bcoo(key, shape, *, dtype=jnp.float_, indices_dtype=None, nse=0.2, n_batch=0, n_dense=0, unique_indices=True, sorted_indices=False, generator=random.uniform, **kwds): """Generate a random BCOO matrix. Args: key : PRNG key to be passed to ``generator`` function. shape : tuple specifying the shape of the array to be generated. dtype : dtype of the array to be generated. indices_dtype: dtype of the BCOO indices. nse : number of specified elements in the matrix, or if 0 < nse < 1, a fraction of sparse dimensions to be specified (default: 0.2). n_batch : number of batch dimensions. must satisfy ``n_batch >= 0`` and ``n_batch + n_dense <= len(shape)``. n_dense : number of batch dimensions. must satisfy ``n_dense >= 0`` and ``n_batch + n_dense <= len(shape)``. unique_indices : boolean specifying whether indices should be unique (default: True). sorted_indices : boolean specifying whether indices should be row-sorted in lexicographical order (default: False). generator : function for generating random values accepting a key, shape, and dtype. It defaults to :func:`jax.random.uniform`, and may be any function with a similar signature. **kwds : additional keyword arguments to pass to ``generator``. Returns: arr : a sparse.BCOO array with the specified properties. """ shape = tuple(map(operator.index, shape)) n_batch = operator.index(n_batch) n_dense = operator.index(n_dense) if n_batch < 0 or n_dense < 0 or n_batch + n_dense > len(shape): raise ValueError(f"Invalid {n_batch=}, {n_dense=} for {shape=}") n_sparse = len(shape) - n_batch - n_dense batch_shape, sparse_shape, dense_shape = map(tuple, split_list(shape, [n_batch, n_sparse])) batch_size = math.prod(batch_shape) sparse_size = math.prod(sparse_shape) if not 0 <= nse < sparse_size: raise ValueError(f"got {nse=}, expected to be between 0 and {sparse_size}") if 0 < nse < 1: nse = int(math.ceil(nse * sparse_size)) nse = operator.index(nse) data_shape = batch_shape + (nse,) + dense_shape indices_shape = batch_shape + (nse, n_sparse) if indices_dtype is None: indices_dtype = dtypes.canonicalize_dtype(jnp.int_) if sparse_size > jnp.iinfo(indices_dtype).max: raise ValueError(f"{indices_dtype=} does not have enough range to generate " f"sparse indices of size {sparse_size}.") @vmap def _indices(key): if not sparse_shape: return jnp.empty((nse, n_sparse), dtype=indices_dtype) flat_ind = random.choice(key, sparse_size, shape=(nse,), replace=not unique_indices).astype(indices_dtype) return jnp.column_stack(jnp.unravel_index(flat_ind, sparse_shape)) keys = random.split(key, batch_size + 1) data_key, index_keys = keys[0], keys[1:] data = generator(data_key, shape=data_shape, dtype=dtype, **kwds) indices = _indices(index_keys).reshape(indices_shape) mat = sparse.BCOO((data, indices), shape=shape) return mat.sort_indices() if sorted_indices else mat