SciJitClass

Callable instances for numba jitclasses. Originally created for SciJIT.

Reference every feature  ·  Guide the long form, with measurements

@scijitclass([('a', float64)])
class Scale:
    def __init__(self, a):
        self.a = a
    def ev_one(self, x):        # stands in for __call__, for a scalar argument
        return self.a * x
    def ev(self, xs):           # stands in for __call__, for an array argument
        return self.a * xs

s = Scale(2.0)
s(3.0)                          # 6.0

Calling a jitclass instance

numba compiles a function marked @njit to machine code, so running that function in a loop does not go through the interpreter. Compiled code accepts a fixed set of types, and an ordinary Python class is not one of them. numba offers @jitclass instead: it compiles a class, and its instances work inside @njit. Every field, meaning every value an instance stores, has to be declared with its type in the decorator.

@jitclass has no equivalent of __call__.

scijitclass registers __call__ through numba's public extension API. The compiler picks the method from the argument types while the calling function compiles.

Install

pip install scijitclass

From a clone:

pip install .

Requires Python 3.10 or later, numba >= 0.66 and numpy. Every number quoted here and in docs/ was measured on numba 0.66.0, numpy 2.4.6, Python 3.14.6. Timings vary with the machine.

Simple example

import numpy as np
from numba import njit, float64
from scijitclass import scijitclass

@scijitclass([('a', float64)])       # field types, exactly as jitclass takes them
class Scale:
    def __init__(self, a):
        self.a = a

    # Scale defines no __call__. These two are what a call forwards to,
    # picked from the type of the argument:

    def ev_one(self, x):             # s(3.0)     one scalar
        return self.a * x

    def ev(self, xs):                # s(array)   one array
        return self.a * xs

s = Scale(2.0)

s(3.0)                               # 6.0
s(np.array([1.0, 2.0]))              # array([2., 4.])

@njit
def use(o, x):
    return o(x)                      # the same call, inside compiled code

use(s, 3.0)                          # 6.0

# A jitclass constructor cannot carry a default inside @njit. A plain @njit
# function can, so the default goes on a one-line factory instead.

@njit
def scale_default(a=2.0):
    return Scale(a)

sd = scale_default()                 # a = 2.0, from the default
sd(3.0)                              # 6.0

@njit
def shifted(x, y):
    o = scale_default(3.0)           # instance built inside compiled code with non-default value
    return o(x) + y

shifted(3, 5.2)                                      # 14.2                -> ev_one
shifted(3, np.array([5.2, 2.5]))                     # array([14.2, 11.5]) -> ev_one
shifted(np.array([3.0, 1.5]), np.array([5.2, 2.5]))  # array([14.2,  7. ]) -> ev

An ordinary Python class defines __call__ and branches inside it on the argument. A jitclass cannot define __call__, so scijitclass puts the branch outside the class, in a table. Each entry names a method and the calls that method accepts.

With no dispatch= argument the table comes from the two method names, and reads:

from scijitclass import all_scalar, first_array

@scijitclass([('a', float64)], dispatch=[('ev', first_array),
                                         ('ev_one', all_scalar)])

Adding more call argument types

A table can name any method and accept any argument type numba can type. One object, four call shapes; the arrows show which method each call runs.

from scijitclass import scijitclass, sig, Scalar, Array, String

@scijitclass([('a', float64)], dispatch=[   # what line(...) forwards to
    ('total',   sig()),              # line()          no arguments
    ('at',      sig(Scalar)),        # line(2.0)       one number
    ('over',    sig(Array)),         # line(array)     one array
    ('by_name', sig(String)),        # line("slope")   one string
])
class Line:
    def __init__(self, a):
        self.a = a

    def total(self):
        return self.a * 100.0

    def at(self, x):
        return self.a * x

    def over(self, xs):
        return self.a * xs

    def by_name(self, what):
        return self.a if what == "slope" else -1.0

line = Line(2.0)

line()                               # 200.0             -> total()
line(3.0)                            # 6.0               -> at()
line(np.array([1.0, 2.0]))           # array([2., 4.])   -> over()
line("slope")                        # 2.0               -> by_name()

Strings, tuples and other jitclass instances select a method the same way.

Usage warnings

Two guards accepting the same call is an error. A guard is the rule saying which calls a method takes. An entry written without one accepts everything numba accepts, and numba is permissive: a method written for an array usually type-checks for a number too. Scalar also covers Integer and Float, so sig(Scalar, Scalar) and sig(Float, Integer) both accept (2.0, 3). That overlap raises AmbiguousDispatch, and the message names both methods.

A dispatch table replaces the ev_one / ev default; it does not extend it. A class defining ev_one that passes dispatch=[('named', ...)] accepts a string and rejects a number. ev_one stays an ordinary method, and obj(...) cannot reach it.

A Python list is not an array. numba gives a Python list its own type, a reflected list, which is not a numpy array and which no array guard admits. obj([1.0, 2.0]) raises TypeError: Scale has no method accepting (reflected list(float64)<iv=None>). Pass np.array(...).

Constructor defaults apply from Python only. A defaulted argument works in the interpreter and raises inside @njit, where every argument is required. Plain jitclasses behave the same way. An @njit factory carries a default in, as scale_default does above.

cache=True never hits on a function taking a jitclass. numba accepts the flag and recompiles every session.

A call from Python costs about 2.5 microseconds: argument conversion, then crossing the box numba wraps an instance in. Inside @njit the compiler picks the method while the calling function compiles, and the call costs what naming the method costs. Build the object once and use it in compiled code. A plain Python class is faster if the object never gets there and the method body is small.

Documentation

docs/REFERENCE.md lists every feature: field types, how a call is declared, the guard catalogue, introspection, costs and restrictions. Read this to look something up.

docs/GUIDE.md is the long form: what a jitclass is, working with several objects at once, and the measurements behind the numbers. Read this to understand why something behaves as it does.

The four scripts in examples/ run in order and print what they measure.

Origin

Written for SciJIT, a package of SciPy-equivalent routines callable inside numba @njit code. Its spline, interpolator and distribution classes are jitclasses, and SciPy spells their evaluation obj(x).

Status

Version 0.1.6

Created with the help of Claude.

Not affiliated with the numba project. scijitclass is a separate package that uses numba's public extension API.

Licensed under BSD-3-Clause. See LICENSE.