Ctrl+K

CustomOpByNumba

Contents

CustomOpByNumba#

class brainpy.math.CustomOpByNumba(eval_shape=None, con_compute=None, name=None, batching_translation=None, jvp_translation=None, transpose_translation=None, multiple_results=True)[source]#

Creating a XLA custom call operator with Numba JIT on CPU backend.

Parameters:

  • name (str) – The name of operator.

  • eval_shape (callable) –

    The function to evaluate the shape and dtype of the output according to the input. This function should receive the abstract information of inputs, and return the abstract information of the outputs. For example:

    >>> def eval_shape(inp1_info, inp2_info, inp3_info, ...):
>>>   return out1_info, out2_info

  • con_compute (callable) –

    The function to make the concrete computation. This function receives inputs, and returns outputs. For example:

    >>> def con_compute(inp1, inp2, inp3, ..., out1, out2, ...):
>>>   pass

Contents