提交 d1c5ae27 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Constants are not inputs

上级 58840bae
......@@ -3,7 +3,6 @@ import warnings
from numpy.random import Generator, RandomState
from pytensor.compile.sharedvalue import SharedVariable, shared
from pytensor.graph.basic import Constant
from pytensor.link.basic import JITLinker
......@@ -72,12 +71,7 @@ class JAXLinker(JITLinker):
def jit_compile(self, fn):
import jax
# I suppose we can consider `Constant`s to be "static" according to
# JAX.
static_argnums = [
n for n, i in enumerate(self.fgraph.inputs) if isinstance(i, Constant)
]
return jax.jit(fn, static_argnums=static_argnums)
return jax.jit(fn)
def create_thunk_inputs(self, storage_map):
from pytensor.link.jax.dispatch import jax_typify
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论