Jax.jit всю мою программу или только ее части? ⇐ Python

[Anonymous]

Предположим, у меня есть программа JAX типа

Код: Выделить всё

def f(x: jnp.array) -> jnp.array:
...

def g(x: jnp.array) -> jnp.array:
# use f lots of times
# do other stuff
...

где f и g полностью совместимы с jit, что мне следует jit?

• Код: Выделить всё

  g

  и позвольте JAX/XLA оптимизировать многократное использование f?

• Код: Выделить всё

  f

  , но не g?

• Код: Выделить всё

  f

  и g?

и почему?
Я пытаюсь понять, что делает jit. Например, разница между jit и XlaBuilder.Build, является ли jit скорее частью Python/JAX и необходим ли он при использовании XLA в целом, а также к каким другим вариантам использования XLA применим шаблон jit.

Подробнее здесь: https://stackoverflow.com/questions/776 ... arts-of-it