Pjit — это функция JAX для компиляции и разбиения вычислений между несколькими ускорительными устройствами. Она обеспечивает распределенное выполнение и шардинг массивов без необходимости полностью переписывать исходную функцию.
A JAX function that splits code to run across multiple accelerator chips. The user passes a function to pjit, which returns a function that has the equivalent semantics but is compiled into an XLA computation that runs across multiple devices (such as GPUs or TPU cores). pjit enables users to shard computations without rewriting them by using the SPMD partitioner. As of March 2023, has been merged with . Refer to Distributed arrays and automatic parallelization for more details.