jax.extend.core.Primitive#

class jax.extend.core.Primitive(name)[source]#
Parameters:

name (str)

__init__(name)[source]#
Parameters:

name (str)

Methods

__init__(name)

abstract_eval(*args, **params)

bind(*args, **params)

bind_with_trace(trace, args, params)

def_abstract_eval(abstract_eval)

def_bind_with_trace(bind_with_trace)

def_effectful_abstract_eval(...)

def_effectful_abstract_eval2(abstract_eval)

def_impl(impl)

get_bind_params(params)

impl(*args, **params)

is_high(**params)

Attributes

call_primitive

is_effectful

map_primitive

multiple_results

ref_primitive

skip_canonicalization

name