import jaxfrom jax import custom_jvpfrom jax import laxfrom jax import vmap
def p_tau(z, tau, alpha=1.5):
return(map_row, alpha=
浏览 24提问于2020-12-30得票数 1
回答已采纳
编辑:这里有GitHub问题：https://github.com/google/jax/issues/5190 我正在尝试使用jit优化以下函数： @partial(jit, static_argnumsitem[1]: return uniques 上面的例程在这里使用astype(np.float32))
浏览 20提问于2020-12-15得票数 1
回答已采纳
1回答
我正在索引向量并使用JAX，但我注意到与简单索引数组时相比，numpy慢了很多。例如，考虑在JAX numpy和普通numpy中创建一个基本数组：import numpy as onp %timeit jax_array[435:852]
速度如此之快，但还是比裸皮慢了2000倍。我正在使用</em
浏览 0提问于2021-08-27得票数 0
回答已采纳
1回答
JAX积累
、、、、
在JAX (如jax.lax.scan )中编译累积内存时，处理内存的最佳方法是什么？然而，尝试使用jax.jit将不出所料地导致ConcretizationTypeError。对于此示例，是否有更有效的内存方法，或者是否以某种方式使用了分配的内存？代码进行行分
浏览 8提问于2022-04-28得票数 1
回答已采纳
我正在学习使用JAX，我对jit和vmap的使用有一些疑问，因为我无法通过阅读文档来解决这些问题。@jax.jit def fooBar(x)：返回foo(x) + bar(x)
如果foo()和bar()已经被抛弃了，有什么区别吗？在上面的例子中，我应该做jax.vmap(fooBar)
浏览 3提问于2021-06-25得票数 2
/JIT版本 import jaxfrom jax import grad, jit, vmapdef jax_simps(f, a,b, N): x,w,y)
@
浏览 39提问于2021-09-10得票数 2
1回答
我想使用JAX的jit方法加速下面示例中嵌套的for-循环。但是，编译需要很长时间，编译后的运行时甚至比不使用jit的版本还要慢。import jax.numpy as jnpfrom jax import random
key
浏览 11提问于2022-01-03得票数 1
回答已采纳
2回答
在Jax中专门在CPU上执行功能
、、、、
如何在CPU上具体地执行我的功能？如果我这样做了：我猜my_function()首先在TPU上执行，结果放在在我的代码开始时使用jax.config.update('jax_platform_name', 'cpu')似乎没有任何效果。谢谢!
浏览 15提问于2022-11-22得票数 1
我希望在jit函数中过滤带有条件的jnp.array，并将其累加到全局变量(因此我们必须使用JAX控制流原语)：import jax.numpy as jnpfrom jax import lax
def get_data():
data = jax.rand
浏览 10提问于2022-06-29得票数 1
在JAX的Quickstart教程中，我发现可以使用以下代码行有效地计算可微函数fun的Hessian矩阵：下面是一个最低限度的工作示例：from jax import jit
浏览 4提问于2022-01-03得票数 3
回答已采纳
2回答
我正在使用JAX，我想执行这样的操作def fun(x, index): return x
这不能在jit下执行。有什么方法可以用jax.ops或jax.lax来完成吗？我曾想过使用jax.ops.index_update(x, idx, y)，但我无法找到一种计算y的方法，而不会再次遇到同样的问题。
浏览 11提问于2021-07-17得票数 4
回答已采纳
我正在实现一个变分自动编码器(VAE)使用Jax和亚麻。在培训期间，我取样了一个潜在的代码(从编码器推断的分布，我使用flax.linen.nn模块的组合来实现)。如何使我的变量成为一个NumPy数组，以便尽可能使用抽象跟踪器以计算效率高的方式实现我的模型？下面是我所面临问题的一个最小的工作示例--我相信需要安装mujoco ()才能运行这个问题：import jax.numpy as np
import numpy as onpimport g
浏览 35提问于2022-08-23得票数 -1
我根据使用pip安装了jax和jaxlib但是在我运行了我的项目之后，它显示了我的恐惧：", line 1, in <module> File "...from 'jaxlib' (/home/control&#x
浏览 39提问于2021-12-13得票数 1
2回答
我想要创建的是一个粗糙的“可微分渲染器”，通过用python编写一个着色器和损失函数，然后使用Jax的AD来找到渐变。然后，我们应该能够通过在这个损失梯度上运行梯度下降来反向渲染图像。我已经让它在简单的着色器上工作得相当好，但当我使用布尔表达式时，我遇到了问题。onp.zeros((WIDTH, HEIGHT, CHANNELS)) jit_func = jax.jit(sh.shade我猜这是因为你不能
浏览 11提问于2020-03-06得票数 1
我在摆弄JAX，通过使用jit装饰器，我发现了两个不同的结果。import jaximport jax.scipy.stats as jstats
return jnp.sum(jnp.log(jstats.multivariate_normal.pdf(x, mean=mu, cov=sigma)))
@jax.j
浏览 9提问于2022-08-27得票数 3
1回答
在计算时间方面，是否有可能使JAX仅减少CPU的性能与Numba相比？+= xidef reduce_1d_jax_serial(x): for xi in x:a = np.random.randn(N)
np.add.reduce(a)给出1.99 µs ...reduce_1d_jax_serial(a).item()给出2
浏览 10提问于2022-04-01得票数 2
回答已采纳
如果我们对数组索引执行非整数计算(然后转换为int() )，那么在jit编译的jax代码中，似乎仍然不能将结果用作有效的索引。我们如何解决这个问题？在木星的细胞里运行这个：from jax import jit@jit d = jnp.sqrt(t.size**2)
浏览 9提问于2021-03-16得票数 1
回答已采纳
我使用jax ()对一个神经网络进行编码，并模拟我的输入，我想从泊松分布中生成一个样本数组。考虑到jax的限制，我该怎么做呢？这两种方法都不起作用，因为它们在jax.numpy和jax.scipy.stats中不受支持。因此，基本上，我需要一个替代的解决方案，要么使用jax支持的另一个包，要么通过硬编码泊松函数。from jax import jit, vmapimpo
浏览 1提问于2019-06-26得票数 2
使用jit进行就地更新似乎对计算时间没有任何影响(这比等效的numpy实现多个数量级)。# 86.9 ns ± 1.42 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
a = np.zeros([node_count,node_count])
print("`jax.np` in place updat
浏览 11提问于2022-11-27得票数 2
回答已采纳
1回答
下面有下面的代码，它使用的是一个简单的for循环。我只是想知道是否有办法实现vmap？signal.lfilter(b, a, [1] + [0]*99) impulse_20 = jnp.asarray(impulse_20)
for ind in range(0, len(y)):
y = jax.ops.index_update(y, jax.ops.i
浏览 0提问于2020-11-03得票数 1
回答已采纳
扫码
添加站长 进交流群
领取专属 10元无门槛券
手把手带您无忧上云
相关资讯
热门标签
云服务器
ICP备案
对象存储
云直播
腾讯会议活动推荐
腾讯云开发者
