How to pass namedtuples and array (made hashable) to a function as static arguments?

[srijaniiserprinceton]

This is a direct reference to the amazing answer by @mattjj to the question of how to make arrays hashable static arguments. In my case, we need to pass a bunch of arguments to the function. Some of them are static and some traced. The way we were doing this was to make namedtuples to create sort of a dictionary and pass them. I'm finding it hard to figure out how to pass both a namedtuple and an array (when made hashable) through the same function. I tried the following (note a slight change in the arguments as compared to the answer by @mattjj).

from functools import partial
from collections import namedtuple
from jax import jit
import jax.numpy as jnp

def some_hash_function(x):
  return int(jnp.sum(x))

class HashableArrayWrapper:
  def __init__(self, val):
    self.val = val
  def __hash__(self):
    return some_hash_function(self.val)
  def __eq__(self, other):
    return (isinstance(other, HashableArrayWrapper) and
            jnp.all(jnp.equal(self.val, other.val)))

def gnool_jit(fun, static_array_argnums=()):
  @partial(jit, static_argnums=static_array_argnums)
  def callee(*args):
    args = list(args)
    for i in static_array_argnums:
      args[i] = args[i].val
    return fun(*args)

  def caller(*args):
    args = list(args)
    for i in static_array_argnums:
      args[i] = HashableArrayWrapper(args[i])
    return callee(*args)

  return caller


###


@partial(gnool_jit, static_array_argnums=(0,1,2))
def f(a,x,y):
  print('re-tracing f!')
  return a * (x ** 2 + y ** 2)


@partial(gnool_jit, static_array_argnums=(0,1,2))
def g(nt,x,y):
  print('re-tracing g!')
  a = nt.a
  return a * (x ** 2 + y ** 2)

x = jnp.array([1,2,3])
y = jnp.array([4,5,6])

a = 1

__ = f(a,x,y)
__ = f(a,x,y)

nt = namedtuple('nt','a')
nt = nt(a)

__ = g(nt,x,y)
__ = g(nt,x,y)

I find that the calls to function f works great but the calls to g fails giving me the following error

ValueError: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class '__main__.HashableArrayWrapper'>, <__main__.HashableArrayWrapper object at 0x2b8099eb8730>. The error was:
TypeError: sum requires ndarray or scalar arguments, got <class '__main__.nt'> at position 0.

I think this is because it is trying to index a namedtuple and failing? How do I get around this problem?

P.S: The motivation behind this is that for my problem, I know the sizes of matrices that will be generated apriori. So, in order for the dimensions to NOT be traced values, we are trying to pre-compute the dimensions and pass the array containing the dimensions. This is where we were trying to make the array (containing the dimensions of arrays to be created in the jitted function) as a static-argument.

[PhilipVinc]

The error is because you creating a HashableArrayWrapper of a namedtuple nt, which then later calls sum(nt(1)) which then fails because nt is of course not a valid input to a jax operation

[srijaniiserprinceton]

I do understand that and mentioned the same in my post just after the error message. You see the point I am trying to make is that namedtuples work just fine when passed as a static argument in the usual jax.jit way. In the answer of @mattjj I cited above, he shows a neat way to passing arrays as static arguments too (which is not allowed by default). I tried extending his method to passing not only arrays as arguments but also namedtuples (so that I can avoid passing an extremely long list of scalars and arrays to each function). This is where is fails (there reason being that it is trying to index a named tuple which is not allowed).

I am looking for a solution to this problem. In my mind, it would be ideal if @gnool_jit could treat only the array static arguments as @mattjj did and the rest of the arguments should be handled just like jax.jit does.

Edit 1: A very non-concrete thought ---> I meant passing the arrays as a separate argument like 'static_array_argnums' and passing the rest of the arguments which may contain namedtuples as a separate argument like 'static_argnum' and then the variables which get passed under the former gets treated a gnool_jit defined above and the ones passed as the latter gets treated by jax.jit.

[PhilipVinc]

Your example code is very hard to read, so I'm not sure what you're trying to do.
But I think that you can achieve this by using pytrees and flax dataclasses, which can have static fields?
I particularly like this approach because it allows me not to have to remember which arguments are static and which aren't, as they will all be handled by the dispatcher automatically for me...

See the example below, which is a jax.jitted function taking a dynamic argument (a png seed) and a dict of shapes.
Notice that the shapes themselves are HashableTuples but I marked no argument as static when calling @jax.jit.

Example here: https://gist.github.com/PhilipVinc/6fa91fe2ce8f71b473a010d676ae4386#file-test-py

By the way, you could also use the same principle to pass in your washable arrays as static objects
See here: https://gist.github.com/PhilipVinc/6fa91fe2ce8f71b473a010d676ae4386#file-bonus-py

And then you can combine the two.

[srijaniiserprinceton]

Sorry for the late response. This answer looks great!