ak.jax.register_behavior_class

Defined in awkward.jax on line 61.

ak.jax.register_behavior_class(cls)

Parameters:

cls – behavior class to register with JAX

Register the behavior class with JAX, if JAX integration is enabled. Otherwise, queue the type for subsequent registration when/if JAX is registered.