理解JAX的追踪器与静态工作方式

huangapple
go评论108阅读模式
英文:

Understanding how JAX's tracer vs static work

问题

以下是您提供的代码的翻译：

  1. 我是JAX的新手，试图使用JAX编写一个简单的代码，在某个时候需要使用scipy方法，然后我想取其导数。
  3. 代码无法运行并出现错误。以下是代码和错误信息。我多次阅读了JAX的文档，但无法弄清楚如何正确编写代码。
  5. import jax.numpy as jnp
  6. from jax import grad, jit, vmap
  7. from jax import random
  8. import numpy as np
  9. import scipy
  10. key = random.PRNGKey(1)
  12. size = 3
  13. x = random.uniform(key, (size, size), dtype=jnp.float32)
  15. def error_func(x):
  16.     dists = scipy.spatial.distance.cdist(x, x, metric='euclidean')
  17.     error = jnp.sum(jnp.array(dists))
  18.     return error
  20. error_diff = grad(error_func)
  22. print(error_func(x))
  23. print(error_diff(x))
  25. 我得到以下错误信息：
  27. 未找到GPU/TPU，回退到CPU。（设置TF_CPP_MIN_LOG_LEVEL=0并重新运行以获取更多信息。）
  28. 3.2158318
  29. Traceback (most recent call last):
  30.   File "/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py", line 26, in <module>
  31.     print(error_diff(x))
  32.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
  33.     return fun(*args, **kwargs)
  34.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/api.py", line 646, in grad_f
  35.     _, g = value_and_grad_f(*args, **kwargs)
  36.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
  37.     return fun(*args, **kwargs)
  38.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/api.py", line 722, in value_and_grad_f
  39.     ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
  40.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/api.py", line 2179, in _vjp
  41.     out_primal, out_vjp = ad.vjp(
  42.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/interpreters/ad.py", line 139, in vjp
  43.     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  44.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/interpreters/ad.py", line 128, in linearize
  45.     jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
  46.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
  47.     return func(*args, **kwargs)
  48.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/interpreters/partial_eval.py", line 777, in trace_to_jaxpr_nounits
  49.     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  50.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
  51.     ans = self.f(*args, **dict(self.params, **kwargs))
  52.   File "/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py", line 15, in error_func
  53.     dists = scipy.spatial.distance.cdist(x, x, metric='euclidean')
  54.   File "/home/sattarian/.local/lib/python3.9/site-packages/scipy/spatial/distance.py", line 2909, in cdist
  55.     XA = np.asarray(XA)
  56.   File "/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/core.py", line 598, in __array__
  57.     raise TracerArrayConversionError(self)
  58. jax._src.traceback_util.UnfilteredStackTrace: jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ConcreteArray([[0.7551559  0.3129729  0.12388372]
  59.  [0.548188   0.8851279  0.30576992]
  60.  [0.82008433 0.95633745 0.3566252 ]], dtype=float32)>with<JVPTrace(level=2/0)> with
  61.   primal = Array([[0.7551559 , 0.3129729 , 0.12388372],
  62.        [0.548188  , 0.8851279 , 0.30576992],
  63.        [0.82008433, 0.95633745, 0.3566252 ]], dtype=float32)
  64.   tangent = Traced<ShapedArray(float32[3,3])>with<JaxprTrace(level=1/0)> with
  65.     pval = (ShapedArray(float32[3,3]), None)
  66.     recipe = LambdaBinding()
  67. See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
  69. The stack trace below excludes JAX-internal frames.
  70. The preceding is the original exception that occurred, unmodified.
  72. --------------------
  74. 上述异常是以下异常的直接原因：
  76. Traceback (most recent call last):
  77.   File "/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py", line 26, in <module>
  78.     print(error_diff(x))
  79.   File "/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py", line 15, in error_func
  80.     dists = scipy.spatial.distance.cdist(x, x, metric='euclidean')
  81.   File "/home/sattarian/.local/lib/python3.9/site-packages/scipy/spatial/distance.py", line 2909, in cdist
  82.     XA = np.asarray(XA)
  83. jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ConcreteArray([[0.7551559  0.3129729  0.12388372]
  84.  [0.548188   0.8851279  0.30576992]
  85.  [0.82008433 0.95633745 0.3566252 ]], dtype=float32)>with<JVPTrace(level=2/0)> with
  86.   primal = Array([[0.7551559 , 0.3129729 , 0.12388372],
  87.        [0.548188  , 
  89. <details>
  90. <summary>英文:</summary>
  92. I&#39;m new to JAX and trying to write a simple code using JAX where at some point I need to use a scipy method. Then I want to take its derivative. 
  95. The code doesn&#39;t run and gives me error. The following is the code and the error. I read a the documentation of JAX a couple of times but couldn&#39;t Figure out what to do to write the code correctly
  96. ```python
  97. import jax.numpy as jnp
  98. from jax import grad, jit, vmap
  99. from jax import random
  100. import numpy as np
  101. import scipy
  102. key = random.PRNGKey(1)
  104. size = 3
  105. x = random.uniform(key, (size, size), dtype=jnp.float32)
  107. def error_func(x):
  108.     dists = scipy.spatial.distance.cdist(x, x, metric=&#39;euclidean&#39;)
  109.     error = jnp.sum(jnp.array(dists))
  110.     return error
  112. error_diff = grad(error_func)
  114. print(error_func(x))
  115. print(error_diff(x))

And I get the followig error:

  1. No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
  2. 3.2158318
  3. Traceback (most recent call last):
  4.   File &quot;/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py&quot;, line 26, in &lt;module&gt;
  5.     print(error_diff(x))
  6.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/traceback_util.py&quot;, line 166, in reraise_with_filtered_traceback
  7.     return fun(*args, **kwargs)
  8.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/api.py&quot;, line 646, in grad_f
  9.     _, g = value_and_grad_f(*args, **kwargs)
  10.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/traceback_util.py&quot;, line 166, in reraise_with_filtered_traceback
  11.     return fun(*args, **kwargs)
  12.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/api.py&quot;, line 722, in value_and_grad_f
  13.     ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
  14.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/api.py&quot;, line 2179, in _vjp
  15.     out_primal, out_vjp = ad.vjp(
  16.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/interpreters/ad.py&quot;, line 139, in vjp
  17.     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  18.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/interpreters/ad.py&quot;, line 128, in linearize
  19.     jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
  20.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/profiler.py&quot;, line 314, in wrapper
  21.     return func(*args, **kwargs)
  22.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/interpreters/partial_eval.py&quot;, line 777, in trace_to_jaxpr_nounits
  23.     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  24.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/linear_util.py&quot;, line 188, in call_wrapped
  25.     ans = self.f(*args, **dict(self.params, **kwargs))
  26.   File &quot;/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py&quot;, line 15, in error_func
  27.     dists = scipy.spatial.distance.cdist(x, x, metric=&#39;euclidean&#39;)
  28.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/scipy/spatial/distance.py&quot;, line 2909, in cdist
  29.     XA = np.asarray(XA)
  30.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/jax/_src/core.py&quot;, line 598, in __array__
  31.     raise TracerArrayConversionError(self)
  32. jax._src.traceback_util.UnfilteredStackTrace: jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced&lt;ConcreteArray([[0.7551559  0.3129729  0.12388372]
  33.  [0.548188   0.8851279  0.30576992]
  34.  [0.82008433 0.95633745 0.3566252 ]], dtype=float32)&gt;with&lt;JVPTrace(level=2/0)&gt; with
  35.   primal = Array([[0.7551559 , 0.3129729 , 0.12388372],
  36.        [0.548188  , 0.8851279 , 0.30576992],
  37.        [0.82008433, 0.95633745, 0.3566252 ]], dtype=float32)
  38.   tangent = Traced&lt;ShapedArray(float32[3,3])&gt;with&lt;JaxprTrace(level=1/0)&gt; with
  39.     pval = (ShapedArray(float32[3,3]), None)
  40.     recipe = LambdaBinding()
  41. See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
  43. The stack trace below excludes JAX-internal frames.
  44. The preceding is the original exception that occurred, unmodified.
  46. --------------------
  48. The above exception was the direct cause of the following exception:
  50. Traceback (most recent call last):
  51.   File &quot;/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py&quot;, line 26, in &lt;module&gt;
  52.     print(error_diff(x))
  53.   File &quot;/mnt/d/OneDrive - UW-Madison/opttest/jaxtest6.py&quot;, line 15, in error_func
  54.     dists = scipy.spatial.distance.cdist(x, x, metric=&#39;euclidean&#39;)
  55.   File &quot;/home/sattarian/.local/lib/python3.9/site-packages/scipy/spatial/distance.py&quot;, line 2909, in cdist
  56.     XA = np.asarray(XA)
  57. jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced&lt;ConcreteArray([[0.7551559  0.3129729  0.12388372]
  58.  [0.548188   0.8851279  0.30576992]
  59.  [0.82008433 0.95633745 0.3566252 ]], dtype=float32)&gt;with&lt;JVPTrace(level=2/0)&gt; with
  60.   primal = Array([[0.7551559 , 0.3129729 , 0.12388372],
  61.        [0.548188  , 0.8851279 , 0.30576992],
  62.        [0.82008433, 0.95633745, 0.3566252 ]], dtype=float32)
  63.   tangent = Traced&lt;ShapedArray(float32[3,3])&gt;with&lt;JaxprTrace(level=1/0)&gt; with
  64.     pval = (ShapedArray(float32[3,3]), None)
  65.     recipe = LambdaBinding()
  66. See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

答案1

得分: 0

JAX转换只适用于JAX函数，而不适用于numpy或scipy函数（这在上面的错误消息中链接中简要讨论）。如果您想要使用grad和其他JAX转换，您需要使用JAX操作编写您的逻辑，而不是来自numpy、scipy或其他不兼容JAX的库的操作。

JAX当前不包括scipy.spatial.distance的任何包装器（尽管有一些正在进行中，参见#16147），因此最好的选择是自己编写代码。幸运的是，cdist相当简单：

  1. def cdist(x, y, metric='euclidean'):
  2.   assert x.ndim == y.ndim == 2
  3.   if metric != 'euclidean':
  4.     raise NotImplementedError(f"{metric=}")
  5.   return jnp.sqrt(jnp.sum((x[:, None, :] - y[None, :, :]) ** 2, axis=-1))
  7. def error_func(x):
  8.     dists = cdist(x, x, metric='euclidean')
  9.     error = jnp.sum(dists)
  10.     return error
  12. error_diff = grad(error_func)
  14. print(error_func(x))
  15. # 3.2158318
  17. print(error_diff(x))
  18. # [[nan nan nan]
  19. #  [nan nan nan]
  20. #  [nan nan nan]]

您会注意到梯度处处为nan。这是预期的结果，因为grad(jnp.sqrt)(0.0)发散（返回无穷大），而根据定义，0.0 * inf等于nan。

英文:

JAX transformations only work on JAX functions, not numpy or scipy functions (this is discussed briefly at the link shown in the error message above) If you want to use grad and other JAX transformations, you need to write your logic using JAX operations, not operations from numpy, scipy, or other non-JAX-compatible libraries.

JAX does not currently include any wrappers of scipy.spatial.distance (though there are some in progress, see #16147), so the best option would be to write the code yourself. Fortunately, cdist is pretty straightforward:

  1. def cdist(x, y, metric=&#39;euclidean&#39;):
  2.   assert x.ndim == y.ndim == 2
  3.   if metric != &#39;euclidean&#39;:
  4.     raise NotImplementedError(f&quot;{metric=}&quot;)
  5.   return jnp.sqrt(jnp.sum((x[:, None, :] - y[None, :, :]) ** 2, axis=-1))
  7. def error_func(x):
  8.     dists = cdist(x, x, metric=&#39;euclidean&#39;)
  9.     error = jnp.sum(dists)
  10.     return error
  12. error_diff = grad(error_func)
  14. print(error_func(x))
  15. # 3.2158318
  17. print(error_diff(x))
  18. # [[nan nan nan]
  19. #  [nan nan nan]
  20. #  [nan nan nan]]

You'll notice that the gradient is everywhere nan. This is the expected result due to the fact that grad(jnp.sqrt)(0.0) diverges (returns infinity), and 0.0 * inf by definition is nan.

huangapple
  • 本文由 发表于 2023年6月13日 06:56:28
  • 转载请务必保留本文链接：https://go.coder-hub.com/76460770.html
  • jax
  • python
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定