brainpy.math module

Basis for Object-oriented Transformations

BrainPyObject([name])	The BrainPyObject class for whole BrainPy ecosystem.
FunAsObject(target[, child_objs, dyn_vars, name])	Transform a Python function as a BrainPyObject.
dyn_seq([iterable])	A list to represent a dynamically changed numerical sequence in which its element can be changed during JIT compilation.
dyn_dict	A dict to represent a dynamically changed numerical dictionary in which its element can be changed during JIT compilation.
Variable(value_or_size[, dtype, batch_axis])	The pointer to specify the dynamical variable.
Parameter(value_or_size[, dtype, batch_axis])	The pointer to specify the parameter.
TrainVar(value_or_size[, dtype, batch_axis])	The pointer to specify the trainable variable.
Partial(fun, *args[, child_objs, dyn_vars])

Object-oriented Transformations

grad(func[, grad_vars, dyn_vars, ...])	Automatic gradient computation for functions or class objects.
vector_grad(func[, grad_vars, dyn_vars, ...])	Take vector-valued gradients for function func.
jacobian(func[, grad_vars, dyn_vars, ...])	Extending automatic Jacobian (reverse-mode) of func to classes.
jacrev(func[, grad_vars, dyn_vars, ...])	Extending automatic Jacobian (reverse-mode) of func to classes.
jacfwd(func[, grad_vars, dyn_vars, ...])	Extending automatic Jacobian (forward-mode) of func to classes.
hessian(func[, grad_vars, dyn_vars, ...])	Hessian of func as a dense array.
make_loop(body_fun, dyn_vars[, out_vars, ...])	Make a for-loop function, which iterate over inputs.
make_while(cond_fun, body_fun, dyn_vars)	Make a while-loop function.
make_cond(true_fun, false_fun[, dyn_vars])	Make a condition (if-else) function.
cond(pred, true_fun, false_fun, operands[, ...])	Simple conditional statement (if-else) with instance of Variable.
ifelse(conditions, branches[, operands, ...])	If-else control flows looks like native Pythonic programming.
for_loop(body_fun, operands[, dyn_vars, ...])	for-loop control flow with Variable.
while_loop(body_fun, cond_fun, operands[, ...])	while-loop control flow with Variable.
to_object([f, child_objs, dyn_vars, name])	Transform a Python function to BrainPyObject.
to_dynsys([f, child_objs, dyn_vars, name])	Transform a Python function to a DynamicalSystem.
function([f, nodes, dyn_vars, name])	Transform a Python function into a BrainPyObject.
jit(func[, dyn_vars, child_objs, ...])	JIT (Just-In-Time) compilation for class objects.
ObjectTransform([name])	Object-oriented JAX transformation for BrainPy computation.

Brain Dynamics Dedicated Operators

pre2post_sum(pre_values, post_num, post_ids)	The pre-to-post synaptic summation.
pre2post_prod(pre_values, post_num, post_ids)	The pre-to-post synaptic production.
pre2post_max(pre_values, post_num, post_ids)	The pre-to-post synaptic maximization.
pre2post_min(pre_values, post_num, post_ids)	The pre-to-post synaptic minimization.
pre2post_mean(pre_values, post_num, post_ids)	The pre-to-post synaptic mean computation.
pre2post_event_sum(events, pre2post, post_num)	The pre-to-post event-driven synaptic summation with CSR synapse structure.
pre2post_coo_event_sum(events, pre_ids, ...)	The pre-to-post synaptic computation with event-driven summation.
pre2post_event_prod(events, pre2post, post_num)	The pre-to-post synaptic computation with event-driven production.
pre2syn(pre_values, pre_ids)	The pre-to-syn computation.
syn2post_sum(syn_values, post_ids, post_num)	The syn-to-post summation computation.
syn2post(syn_values, post_ids, post_num[, ...])	The syn-to-post summation computation.
syn2post_prod(syn_values, post_ids, post_num)	The syn-to-post product computation.
syn2post_max(syn_values, post_ids, post_num)	The syn-to-post maximum computation.
syn2post_min(syn_values, post_ids, post_num)	The syn-to-post minimization computation.
syn2post_mean(syn_values, post_ids, post_num)	The syn-to-post mean computation.
syn2post_softmax(syn_values, post_ids, post_num)	The syn-to-post softmax computation.
sparse_matmul(A, B)	Sparse matrix multiplication.
csr_matvec(values, indices, indptr, vector, ...)	Product of CSR sparse matrix and a dense vector.
event_csr_matvec(values, indices, indptr, ...)	The pre-to-post event-driven synaptic summation with CSR synapse structure.
segment_sum(data, segment_ids[, ...])	segment_sum operator for brainpy Array and Variable.
segment_prod(data, segment_ids[, ...])	segment_prod operator for brainpy Array and Variable.
segment_max(data, segment_ids[, ...])	segment_max operator for brainpy Array and Variable.
segment_min(data, segment_ids[, ...])	segment_min operator for brainpy Array and Variable.
XLACustomOp([eval_shape, con_compute, ...])	Creating a XLA custom call operator.

Activation Functions

celu(x[, alpha])	Continuously-differentiable exponential linear unit activation.
elu(x[, alpha])	Exponential linear unit activation function.
gelu(x[, approximate])	Gaussian error linear unit activation function.
glu(x[, axis])	Gated linear unit activation function.
hard_tanh(x)	Hard \(\mathrm{tanh}\) activation function.
hard_sigmoid(x)	Hard Sigmoid activation function.
hard_silu(x)	Hard SiLU activation function
hard_swish(x)	Hard SiLU activation function
leaky_relu(x[, negative_slope])	Leaky rectified linear unit activation function.
log_sigmoid(x)	Log-sigmoid activation function.
log_softmax(x[, axis])	Log-Softmax function.
one_hot(x, num_classes, *[, dtype, axis])	One-hot encodes the given indicies.
normalize(x[, axis, mean, variance, epsilon])	Normalizes an array by subtracting mean and dividing by sqrt(var).
relu(x)
relu6(x)	Rectified Linear Unit 6 activation function.
sigmoid(x)	Sigmoid activation function.
soft_sign(x)	Soft-sign activation function.
softmax(x[, axis])	Softmax function.
softplus(x)	Softplus activation function.
silu(x)	SiLU activation function.
swish(x)	SiLU activation function.
selu(x)	Scaled exponential linear unit activation.
identity(n[, dtype])	Return the identity array.
tanh(x)	Compute hyperbolic tangent element-wise.

Array Operations

flatten(input[, start_dim, end_dim])	Flattens input by reshaping it into a one-dimensional tensor.
fill_diagonal(a, val)
remove_diag(arr)	Remove the diagonal of the matrix.
clip_by_norm(t, clip_norm[, axis])
empty(shape[, dtype])	Return a new array of given shape and type, without initializing entries.
empty_like(prototype[, dtype, shape])	Return a new array with the same shape and type as a given array.
ones(shape[, dtype])	Return a new array of given shape and type, filled with ones.
ones_like(a[, dtype, shape])	Return an array of ones with the same shape and type as a given array.
zeros(shape[, dtype])	Return a new array of given shape and type, filled with zeros.
zeros_like(a[, dtype, shape])	Return an array of zeros with the same shape and type as a given array.
array(object[, dtype, copy, order, ndmin])	Create an array.
asarray(a[, dtype, order])	Convert the input to an array.
arange(start[, stop, step, dtype])	Return evenly spaced values within a given interval.
linspace(start, stop[, num, endpoint, ...])	Return evenly spaced numbers over a specified interval.
logspace(start, stop[, num, endpoint, base, ...])	Return numbers spaced evenly on a log scale.
as_device_array(tensor[, dtype])	Convert the input to a jax.numpy.DeviceArray.
as_jax(tensor[, dtype])	Convert the input to a jax.numpy.DeviceArray.
as_ndarray(tensor[, dtype])	Convert the input to a numpy.ndarray.
as_numpy(tensor[, dtype])	Convert the input to a numpy.ndarray.
as_variable(tensor[, dtype])	Convert the input to a brainpy.math.Variable.

Delay Variables

TimeDelay(delay_target, delay_len[, ...])	Delay variable which has a fixed delay time length.
LengthDelay(delay_target, delay_len[, ...])	Delay variable which has a fixed delay length.
NeuTimeDelay(delay_target, delay_len[, ...])	Neutral Time Delay.
NeuLenDelay(delay_target, delay_len[, ...])	Neutral Length Delay.

Environment Settings

set_float(dtype)	Set global default float type.
get_float()	Get the default float data type.
set_int(dtype)	Set global default integer type.
get_int()	Get the default int data type.
set_bool(dtype)	Set global default boolean type.
get_bool()	Get the default boolean data type.
set_complex(dtype)	Set global default complex type.
get_complex()	Get the default complex data type.
set_dt(dt)	Set the default numerical integrator precision.
get_dt()	Get the numerical integrator precision.
set_mode(mode)	Set the default computing mode.
get_mode()	Get the default computing mode.
set_environment([mode, dt, x64, complex_, ...])	Set the default computation environment.
enable_x64()
disable_x64()
set_platform(platform)	Changes platform to CPU, GPU, or TPU.
get_platform()	Get the computing platform.
set_host_device_count(n)	By default, XLA considers all CPU cores as one device.
clear_buffer_memory([platform])	Clear all on-device buffers.
enable_gpu_memory_preallocation()	Disable pre-allocating the GPU memory.
disable_gpu_memory_preallocation()	Disable pre-allocating the GPU memory.
ditype()	Default int type.
dftype()	Default float type.
environment([mode, dt, x64, complex_, ...])	Context-manager that sets a computing environment for brain dynamics computation.
batching_environment([dt, x64, complex_, ...])	Environment with the batching mode.
training_environment([dt, x64, complex_, ...])	Environment with the training mode.

Computing Modes

Mode()	Base class for computation Mode
NonBatchingMode()	Normal non-batching mode.
BatchingMode()	Batching mode.
TrainingMode()	Training mode requires data batching.
nonbatching_mode	Normal non-batching mode.
batching_mode	Batching mode.
training_mode	Training mode requires data batching.