1回答
import jax.numpy as jnp return jnp.sqrt为了计算xt、yt和zt中每对数据点的梯度，我必须执行以下操作： fn_prime = vmap(grad(fn, argnums=(0, 1, 2)), in_axes=(None, None, None
浏览 22提问于2021-10-17得票数 1
1回答
矢量嵌套vmap
、、
以下是我掌握的一些数据：import numpyro.distributions as distjax.vmap(lambda axis: jax.vmap(func, (None, 0))(axis, yaxis))(xaxis)另一种方
浏览 5提问于2021-11-04得票数 2
回答已采纳
我正在尝试理解jax.vmap/pmap (jax：https://jax.readthedocs.io/)的行为。CUDA有threadId让你知道哪个线程在执行代码，在jax中有类似的概念吗？(jax.process_id没有)
浏览 61提问于2021-10-05得票数 1
回答已采纳
2回答
我想从一个函数构造一个2D数组，这样我就可以利用jax.jit了。 我通常使用numpy的方法是创建一个空数组，然后就地填充该数组。xx[ii_a, ii_b] = xyz[0] zz[ii_a, ii_b] = xyz[2] 为了在jax中实现这一点，我尝试使用jax.opt.index_update。yy = yy.at[ii_a, ii_b].set(xyz[1])
zz = zz.at[ii_a, ii_b].set(xyz[2]) 它运
浏览 33提问于2020-10-06得票数 1
回答已采纳
1回答
为了提高性能，我想使用vmap来向量化这段代码。b in dataB] for a in dataA]) return func(x, y)
mapped = jax.vmap
浏览 24提问于2021-10-04得票数 3
回答已采纳
我正在尝试理解vmap中的in_axes和out_axes选项。例如，我想要将两个矩阵相加，得到具有相同形状的输出。X = np.arange(9).reshape(3,3)def sum2(x,y):vmapipython-input-403-103694166574> in <module> 4 return x + y
----> 5 vmap</
浏览 33提问于2021-07-11得票数 2
回答已采纳
输入/ (现在我将坚持使用Autograd，它提供了一个autograd.jacobian()方法，但只要得到我想要的东西，我就很乐意使用JAX )。**( )(
>>> import autograd.numpy as np>>> x = np.array(
浏览 5提问于2020-01-14得票数 2
回答已采纳
2回答
., dtype=float32)]jax.vmap(f)(zip(xs, ys))ValueError: vmap was requested to map its argument along axis
浏览 9提问于2022-06-05得票数 1
回答已采纳
1回答
我是Jax的新手，我正在努力转换别人的代码，它使用了numba的“快速数学”特性，并且依赖于许多嵌套的for-循环，而不造成太多的性能损失。我正在尝试使用Jax的vmap函数重新创建相同的行为。下面是我试图用vmap向量化的一个简单的例子：from jax import vmap
a = jnp.arange(20(a, jax.ops.index[i,
浏览 3提问于2021-05-23得票数 2
回答已采纳
我正在学习使用JAX，我对jit和vmap的使用有一些疑问，因为我无法通过阅读文档来解决这些问题。@jax.jit def fooBar(x)：返回foo(x) + bar(x)
如果foo()和bar()已经被抛弃了，有什么区别吗？，，在我把它放进vmap之后，我应该设置一个函数吗？在上面的例子中，我应该做jax.vmap(fooBar)还是jax<
浏览 3提问于2021-06-25得票数 2
我正在阅读一套很好的JAX笔记本，其中包含关于的练习。我感兴趣的任务是通过多个vmap函数应用程序复制以下函数： s = vmap(np.sum)(cp)我最初的“猜测解决方案”是尝试：
func1 = <em
浏览 5提问于2022-10-09得票数 1
回答已采纳
1回答
在JAX中，我希望vmap是一个固定长度的数据类型列表上的函数，例如：from flax import struct
classparam_list = jnp.Array([EnvParams(max_steps=500), EnvParams(max_steps=600)])上面的示例失败了，因为无法创建自定义对象的jnp.Array，而且JAX不允许在Py
浏览 15提问于2022-09-18得票数 0
1回答
numpy as np a = x**3+5
g2i = jax.vmap(jax.grad(jax.grad(f)))(xi)
g3i = jax.vmap(jax.grad(jax.grad(<
浏览 3提问于2021-07-30得票数 0
回答已采纳
1回答
我只是想知道是否有办法实现vmap？def filter_jax(y): y = jax.ops.index_update(y, jax.ops.index[(jnpData).block_until_ready()def paraUpdate(y, ind):
浏览 0提问于2020-11-03得票数 1
回答已采纳
在使用JAX进行优化之前，此函数的代码如下： """ Calculate the unweighted importanceresult[ii] = pi_ratio_prod * weighted_return
不幸的是，我不能仅仅重写这个函数来处理一个事件，然后使用jax.vmap原因是我想要矢量化的参数是一个定制的Episode对象，JAX不
浏览 2提问于2022-10-01得票数 0
回答已采纳
1回答
我试图理解JAX vmap的行为，所以我编写了以下代码：from jax import vmap
z = jnp.dot(a,b)b = jnp.array
浏览 10提问于2021-03-09得票数 2
回答已采纳
下面是一个示例：from jax import map
return jnp.dot(A, B)
R = vmap(我希望结果是形状(10, 10, 2, 2)，因为我认为vmap将函数应用于指定轴的每个子数组。显然，我没有正确地理解它是如何
浏览 4提问于2022-04-10得票数 1
我正在尝试使用vmap来理解JAX的自动矢量化功能，并根据JAX的文档实现了一个最小的工作示例。import jax.numpy as jnp
input_batch = jnp.ones((batc
浏览 8提问于2022-01-03得票数 4
回答已采纳
对于自动区分，我使用JAX库。为了计算我的目标变量'u‘的一阶和二阶导数，即计算du/dx和d2u/dx2，在一个例子中使用了elementwise_grad。在jax中，它的替代方案是什么？通过JAX，我应该使用哪些函数来计算du/dx和d2u/dx2？
浏览 4提问于2021-12-08得票数 1
云服务器
ICP备案
对象存储
云点播
实时音视频
腾讯云开发者
Copyright © 2013 - 2026 Tencent Cloud. All Rights Reserved. 腾讯云 版权所有
深圳市腾讯计算机系统有限公司 ICP备案/许可证号：粤B2-20090059 
粤公网安备44030502008569号
腾讯云计算（北京）有限责任公司 京ICP证150476号 | 京ICP备11018762号