Skip to content

Commit

Permalink
Constants are not inputs
Browse files Browse the repository at this point in the history
  • Loading branch information
ricardoV94 committed Nov 29, 2024
1 parent 58840ba commit d1c5ae2
Showing 1 changed file with 1 addition and 7 deletions.
8 changes: 1 addition & 7 deletions pytensor/link/jax/linker.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
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


Expand Down Expand Up @@ -72,12 +71,7 @@ def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
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
Expand Down

0 comments on commit d1c5ae2

Please sign in to comment.