Need a way to specify constants

I have a use case to jit a function that does slicing depending on a dict arguments

@jit
def foo(x, d: dict):
  return x[:d['num']] + d['val']

foo(jnp.arange(10), {'num': 5, 'value': np.arange(5)})

no surprise I got IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax.

static_argnums is yet not working as expected since I have an array in d

@partial(jit, static_argnums=1)
def foo(x, d: dict):
  return x[:d['num']] + d['val']

foo(jnp.arange(10), {'num': 5, 'value': np.arange(5)})

I got ValueError: Non-hashable static arguments are not supported.

Here is a working dirty trick

@jit
def bar(x, d: dict):
  return x[:d['num'].shape[0]] + d['val']

bar(jnp.arange(10), {'num': jnp.zeros((5, 0)), 'val': np.arange(5)})

I got DeviceArray([0, 2, 4, 6, 8], dtype=int32)

Is there a way to specify argument d['num'] as constant?

The above weird use case is more reasonable in Flax whose variables are collected as a dict input.

[NeilGirdhar]

I want to understand better the difference between the two, since I noticed that in the latter case I can actually use non-hashable

If you put non-hashable data into static fields, the JIT lookup will fail. Jax doesn’t give you an error for putting unhashable data in static fields yet.

There’s some more information & discussion of this foot-gun in this issue

Another related issue #7826

[jakevdp]

For any custom class Python defaults to creating a hash function based on the object ID, without reference to any of the classes members or contents:

import jax.numpy as jnp

class Foo:
  def __init__(self, x):
    self.x = x

f = Foo(jnp.arange(4))
first_hash = hash(f)

f.x = jnp.ones(10)
second_hash = hash(f)

print("This is bad if True:", first_hash == second_hash)
# This is bad if True: True

This means that if you create a PyTree class and use it as a static argument, then mutate it, and pass it to the function again, you will get unexpected results. There’s some more information & discussion of this foot-gun in this issue: https://github.com/google/jax/issues/9024

If you create a new PyTree class each time, you will have a different object (and thus a different hash) and your function will be recompiled each time you call it.

There are ways you can hack around using jnp.ndarray as a static argument (essentially by wrapping it in some class that defines __hash__), but I wouldn’t recommend it because it’s easy to inadvertently do the wrong thing, and this kind of problem is usually better solved by re-writing your code in a manner that conforms to typical JAX approaches.