brainpy.math.apply_along_axis#

brainpy.math.apply_along_axis(func1d, axis, arr, *args, **kwargs)#

Similar to jax.numpy.apply_along_axis function, while it is compatible with brainpy Array/Variable.

Note that this function is also compatible with:

  1. NumPy or PyTorch syntax when receiving out argument.

  2. PyTorch syntax when receiving keepdim or dim argument.

  3. TensorFlow syntax when receiving keep_dims argument.