transonic.backends.jax

Jax backend

Internal API

class transonic.backends.jax.SubBackendJITJax(name, type_formatter)[source]

Bases: SubBackendJITPython

class transonic.backends.jax.JaxBackend[source]

Bases: PythonBackend

Main class for the Jax backend

See https://github.com/google/jax

_SubBackendJIT

alias of SubBackendJITJax

_make_backend_code(path_py, analysis, **kwargs)[source]

Create a backend code from a Python file

Functions

add_jax_comments(code)

Add Jax code in Python comments

Classes

JaxBackend()

Main class for the Jax backend

SubBackendJITJax(name, type_formatter)