3

これはこの質問に関連しています。奇妙なことの1つを除いて、私はなんとかコードの作業を最大限に活用しています。

これが変更されたコードです。

import jax.numpy as jnp
from jax import grad, jit, value_and_grad
from jax import vmap, pmap
from jax import random
import jax
from jax import lax
from jax import custom_jvp


def p_tau(z, tau, alpha=1.5):
    return jnp.clip((alpha - 1) * z - tau, a_min=0) ** (1 / (alpha - 1))


def get_tau(tau, tau_max, tau_min, z_value):
    return lax.cond(z_value < 1,
                    lambda _: (tau, tau_min),
                    lambda _: (tau_max, tau),
                    operand=None
                    )


def body(kwargs, x):
    tau_min = kwargs['tau_min']
    tau_max = kwargs['tau_max']
    z = kwargs['z']
    alpha = kwargs['alpha']

    tau = (tau_min + tau_max) / 2
    z_value = p_tau(z, tau, alpha).sum()
    taus = get_tau(tau, tau_max, tau_min, z_value)
    tau_max, tau_min = taus[0], taus[1]
    return {'tau_min': tau_min, 'tau_max': tau_max, 'z': z, 'alpha': alpha}, None

@jax.partial(jax.jit, static_argnums=(2,))
def map_row(z_input, alpha, T):
    z = (alpha - 1) * z_input

    tau_min, tau_max = jnp.min(z) - 1, jnp.max(z) - z.shape[0] ** (1 - alpha)
    result, _ = lax.scan(body, {'tau_min': tau_min, 'tau_max': tau_max, 'z': z, 'alpha': alpha}, xs=None,
                         length=T)
    tau = (result['tau_max'] + result['tau_min']) / 2
    result = p_tau(z, tau, alpha)
    return result / result.sum()

@jax.partial(jax.jit, static_argnums=(1,3,))
def _entmax(input, axis=-1, alpha=1.5, T=20):
    result = vmap(jax.partial(map_row, alpha=alpha, T=T), axis)(input)
    return result

@jax.partial(custom_jvp, nondiff_argnums=(1, 2, 3,))
def entmax(input, axis=-1, alpha=1.5, T=10):
    return _entmax(input, axis, alpha, T)
    
@jax.partial(jax.jit, static_argnums=(0,2,))
def _entmax_jvp_impl(axis, alpha, T, primals, tangents):
    input = primals[0]
    Y = entmax(input, axis, alpha, T)
    gppr = Y  ** (2 - alpha)
    grad_output = tangents[0]
    dX = grad_output * gppr
    q = dX.sum(axis=axis) / gppr.sum(axis=axis)
    q = jnp.expand_dims(q, axis=axis)
    dX -= q * gppr
    return Y, dX


@entmax.defjvp
def entmax_jvp(axis, alpha, T, primals, tangents):
    return _entmax_jvp_impl(axis, alpha, T, primals, tangents)


import numpy as np
input = jnp.array(np.random.randn(64, 10)).block_until_ready()
weight = jnp.array(np.random.randn(64, 10)).block_until_ready()

def toy(input, weight):
    return (weight*entmax(input, axis=0, alpha=1.5, T=20)).sum()

jax.jit(value_and_grad(toy))(input, weight)

このコードは、次のようなエラーを生成します。

tuple index out of range

これは、このコード行が原因です

@jax.partial(jax.jit, static_argnums=(2,))
def map_row(z_input, alpha, T):

関数本体をエンティティ関数だけに置き換えても、エラーは解決しません。これは本当に奇妙な振る舞いです。ただし、ループを展開するのに役立つため、これを静的にすることは非常に重要です。

4

1 に答える 1

2

このエラーは、JAX ですぐに修正されることを願っている疣贅によるものです。静的引数はキーワードで渡すことができません。つまり、これを変更する必要があります。

def toy(input, weight):
    return (weight*entmax(input, axis=0, alpha=1.5, T=20)).sum()

これに：

def toy(input, weight):
    return (weight*entmax(input, 0, 1.5, 20)).sum()

への呼び出しにも同じ修正を適用する必要がありますmax_row。

この時点で、トレースされた変数を静的引数を必要とする関数に渡すため、ValueError が発生します。解決策は、JIT で JAX reshape を処理する方法に似ています。

追加の注意: このstatic_argnumsエラーは最近改善されており、次のリリースではもう少し明確になります。

ValueError: jitted function has static_argnums=(2,), donate_argnums=() but was called with only 1 positional arguments.
于 2021-01-07T14:23:35.610 に答える