Использование Jax Jit для метода в качестве декоратора вместо прямого применения функции jit

Декораторы не имеют ничего общего со статическими аргументами: статические аргументы — это концепция, специфичная для jax.jit.

Выполняя резервное копирование, вы должны иметь в виду, что всякий раз, когда jax.jit компилирует функцию, он кэширует артефакт компиляции на основе нескольких величин, в том числе:

1. идентификатор компилируемой функции или вызываемого объекта
2. статические атрибуты любых нестатических аргументов, таких как shape и dtype
3. хэш любых аргументов, помеченных как статические с помощью static_argnums или static_argnames
4. значение любых глобальных конфигураций, которые могут повлиять на результаты

Имея это в виду, давайте рассмотрим этот фрагмент:

c = CustomClass(2, True)
jitted_calc = jit(c.calc)
print(jitted_calc(3))
c.mul = False 
print(jitted_calc(3))

причина, по которой jitted_calc не обновляется при обновлении атрибутов c, заключается в том, что ничего, связанное с ключом кэша, не изменилось: (1) идентификатор функции тот же, (2) форма и тип аргумента не изменились, (3 ) нет статических аргументов, (4) глобальные конфигурации не изменились. Таким образом, предыдущий кэшированный артефакт компиляции (с предыдущим значением mul) выполняется снова. Это основная причина, по которой я не упомянул эту стратегию в документе, на который вы ссылаетесь: это редко бывает то поведение, которое хотелось бы пользователям.

Этот подход к обертыванию связанного метода в JIT, кстати, похож на обертывание определения метода с помощью @partial(jit, static_argnums=0), но детали не такие: в версии static_argnumsself помечен как статический аргумент, и поэтому его хэш становится частью JIT. кэш. Метод __hash__ по умолчанию для класса просто основан на идентификаторе экземпляра, поэтому изменение c.mul не меняет хеш и не запускает перекомпиляцию. Вы можете увидеть пример того, как это исправить, в разделе Стратегия 2 в документе, на который вы ссылаетесь: по сути, определите соответствующие методы __hash__ и __eq__ для класса:

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  @partial(jit, static_argnums=0)
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

  def __hash__(self):
    return hash((self.x, self.mul))

  def __eq__(self, other):
    return (isinstance(other, CustomClass) and
            (self.x, self.mul) == (other.x, other.mul))

В вашем последнем примере вы определяете это:

def decorator(func):
    def wrapper(*args, **kwargs):
        x = func(*args, **kwargs)
        return x
    return wrapper

Этот код вообще не использует jax.jit. Тот факт, что изменения в c.mul приводят к изменениям в выходных данных, не имеет ничего общего с синтаксисом декоратора, а скорее связан с тем, что здесь не задействован JIT-кеш.

Надеюсь, все понятно!