hawk.primitive

Contents

hawk.primitive#

hawk.primitive(name=None, *, vjp=None, jvp=None)[source]#

Register a custom primitive, named name or after the function.

Used as a decorator on the forward:

@hawk.primitive("softplus", vjp=lambda x, bar: bar / (1.0 + exp(-x)))
def softplus(x):
    return log(1.0 + exp(x))

@hawk.primitive alone names it softplus after the function. The decorated object is a PrimitiveDef, callable from a kernel body like a free function. With neither rule it is a plain inline helper whose derivative goes through its forward.

Registering the same function again (a notebook cell run twice) is a no-op that gives back the first registration. A different function under a taken name is refused: a collision would make two kernels’ walk digests agree while their derivatives disagreed.

Return type:

Callable

Parameters:

  • name (Callable | str | None)

  • vjp (Callable | None)

  • jvp (Callable | None)

Contents