braincore.transform module

This module contains the functions for the transformation of the brain data.

Transform Functions

vector_grad([func, grad_vars, argnums, ...])	Take vector-valued gradients for function func.
grad([fun, grad_vars, argnums, holomorphic, ...])	Compute the gradient of a scalar-valued function with respect to its arguments.
jacrev(func[, grad_vars, argnums, has_aux, ...])	Extending automatic Jacobian (reverse-mode) of func to classes.
jacfwd(func[, grad_vars, argnums, has_aux, ...])	Extending automatic Jacobian (forward-mode) of func to classes.
jacobian(func[, grad_vars, argnums, ...])	Extending automatic Jacobian (reverse-mode) of func to classes.
hessian(func[, grad_vars, argnums, ...])	Hessian of func as a dense array.
cond(pred, true_fun, false_fun, *operands)	Conditionally apply true_fun or false_fun.
switch(index, branches, *operands)	Apply exactly one of branches given by index.
ifelse(conditions, branches, *operands[, ...])	If-else control flows looks like native Pythonic programming.
scan(f, init, xs[, length, reverse, unroll, ...])	Scan a function over leading array axes while carrying along state.
for_loop(f, *xs[, length, reverse, unroll, pbar])	for-loop control flow with State.
while_loop(cond_fun, body_fun, init_val)	Call body_fun repeatedly in a loop while cond_fun is True.
jit([fun, in_shardings, out_shardings, ...])	Sets up fun for just-in-time compilation with XLA.

Transform Tools

StatefulFunction	A wrapper class for a function that collects the states that are read and written by the function.
ProgressBar
make_jaxpr	Creates a function that produces its jaxpr given example args.
jit_error	Check errors in a jit function.