Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
View this page
Toggle table of contents sidebar

# autoray.compiler

## Attributes

`_backend_lookup` |
`_compiler_lookup` |

## Classes

`CompilePython` | A simple compiler that unravels all autoray calls, optionally sharing
`CompileJax` |
`CompileTensorFlow` |
`CompileTorch` |
`CompileTorch2` |
`CompilePytensor` |
`AutoCompiled` | Just in time compile a `autoray.do` using function. See the main

## Functions

`autojit`([fn, backend, compiler_opts]) | Just-in-time compile an `autoray` function, automatically choosing

## Module Contents

class autoray.compiler.CompilePython(fn, fold_constants=True, share_intermediates=True)

  A simple compiler that unravels all autoray calls, optionally sharing intermediates and folding constants, converts this to a code object using `compile`, then executes this using `exec`.

  Parameters:

    - fn (callable) – Function to compile - should have signature `fn(*args, **kwargs) -> array`, with `args` and `kwargs` any nested combination of `tuple`, `list` and `dict` objects containing arrays (or other constant arguments), and perform array operations on these using `autoray.do`.

    - fold_constants (bool, optional) – Whether to fold all constant array operations into the graph, which might increase memory usage.

    - share_intermediates (bool, optional) – Whether to cache all computational nodes during the trace, so that any shared intermediate results can be identified.

  _fn
  _fold_constants
  _share_intermediates
  _jit_fn = None
  setup(args, kwargs)

    Convert the example arrays to lazy variables and trace them through the function.

  __call__(*args, array_backend=None, **kwargs)

    If necessary, build, then call the compiled function.

class autoray.compiler.CompileJax(fn, enable_x64=None, platform_name=None, **kwargs)
  _fn
  _enable_x64
  _platform_name
  _jit_fn = None
  _jit_kwargs
  setup()
  __call__(*args, array_backend=None, **kwargs)

class autoray.compiler.CompileTensorFlow(fn, **kwargs)
  _fn
  _jit_fn = None
  _jit_kwargs
  setup()
  __call__(*args, array_backend=None, **kwargs)

class autoray.compiler.CompileTorch(fn, **kwargs)
  torch
  _fn
  _jit_fn = None
  _jit_kwargs
  setup(*args, **kwargs)
  __call__(*args, array_backend=None, **kwargs)

class autoray.compiler.CompileTorch2(fn, **kwargs)
  torch
  _jit_fn
  __call__(*args, array_backend=None, **kwargs)

class autoray.compiler.CompilePytensor(fn, **kwargs)
  _fn
  _jit_fn = None
  _jit_kwargs
  _output_ref_tree = None
  setup(args, kwargs)
  __call__(*args, array_backend=None, **kwargs)

autoray.compiler._backend_lookup

autoray.compiler._compiler_lookup

class autoray.compiler.AutoCompiled(fn, backend=None, compiler_opts=None)

  Just in time compile a `autoray.do` using function. See the main wrapper `autojit`.

  _fn
  _backend
  _compiled_fns
  __call__(*args, backend=None, **kwargs)

autoray.compiler.autojit(fn=None, *, backend=None, compiler_opts=None)

  Just-in-time compile an `autoray` function, automatically choosing the backend based on the input arrays, or via keyword argument.

  The backend used to do the compilation can be set in three ways:

  > 1. Automatically based on the arrays the function is called with, i.e. `cfn(*torch_arrays)` will use `torch.jit.trace`.

  > 2. In this wrapper, `@autojit(backend='jax')`, to provide a specific default instead.

  > 3. When you call the function `cfn(*arrays, backend='torch')` to override on a per-call basis.

  If the arrays supplied are of a different backend type to the compiler, then the returned array will also be converted back, i.e. `cfn(*numpy_arrays, backend='tensorflow')` will return a `numpy` array.

  The `'python'` backend simply extracts and unravels all the `do` calls into a code object using `compile` which is then run with `exec`. This makes use of shared intermediates and constant folding, strips away any python scaffoliding, and is compatible with any library, but the resulting function is not ‘low-level’ in the same way as the other backends.

  Parameters:

    - fn (callable) – The autoray function to compile.

    - backend (str, optional) –

      If set, use this as the default backend. The options are:

      - `'python'`: extract and unravel all the `do` calls into a code object using `compile` which is then run with `exec`.

      - `'jax'`: use jax.jit to compile the function.

      - `'tensorflow'`: use tf.function to compile the function.

      - `'torch'`: use torch.jit.trace to compile the function.

      - `'pytensor'`: use pytensor.function to compile the function.

      If not set, the backend will be inferred from the input arrays, or it can be set at call time with the `backend` keyword argument.

    - compiler_opts (dict[dict], optional) – Dict of dicts when you can supply options for each compiler backend separately, e.g.: `@autojit(compiler_opts={'tensorflow': {'jit_compile': True}})`.

  Returns:

    cfn – The function with auto compilation.

  Return type:

    callable