multimesh.jax.context

Contents

multimesh.jax.context#

multimesh.jax.context(autoshard=None, enable_recomputation=None, enable_fast_path=None, ignore_transforms=None)#

Helper function to configure multiple contexts in a single manager.

Parameters:
  • autoshard (bool | None) – optional, value to configure the autoshard() context manager.

  • enable_fast_path (bool | None) – optional, value to configure the enable_fast_path() context manager.

  • ignore_transforms (bool | None) – optional, value to configure the ignore_transforms() context manager.

  • enable_recomputation (bool | None)