Using tracers as dictionary keys

[EdanToledo]

Is there any solution to using tracers as dictionary keys and allowing their hash to be equivalent to their runtime values?

For example:

import jax
from typing import Dict

sample_dict = {"agent_0": 0, "agent_1": 1, "agent_2": 2, "agent_3": 3}

def unjitted_fun(example_id: int, dictionary: Dict[str, int]):
    example_key = f"agent_{example_id}"
    example_output = dictionary[example_key]

    return example_output

for i in range(4):
    example_output = unjitted_fun(i, sample_dict)
    print(example_output)

This outputs:

0
1
2
3

When this code is jitted i.e:

jitted_fun = jax.jit(unjitted_fun)

for i in range(4):
    example_output = jitted_fun(i, sample_dict)
    print(example_output)

This doesn't work due to a key error as follows:

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/edan/Desktop/examples/test.py", line 21, in <module>
    example_output = jitted_fun(i, sample_dict)
  File "/home/edan/Desktop/examples/test.py", line 7, in unjitted_fun
    example_output = dictionary[example_key]
KeyError: 'agent_Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>'

Any help would be appreciated.

[jakevdp]

The short answer is no: there is no way to make the hash of a tracer match its runtime value, because a tracer is an abstract representation of all possible runtime values.

I suspect this is an XY problem – if you can give a more detailed example of what you're trying to accomplish, we may be able to recommend a viable approach.

[EdanToledo]

I realise now that this is possible and it solves our problem.

[jakevdp]

The jit cache itself might take care of this for you; for example:

@partial(jit, static_argnames=['key'])
def f(x, key):
  print('tracing')
  if key == 0:
    return x
  elif key == 1:
    return [x, x]

result1 = f(0, 0)  # prints "tracing"
result2 = f(0, 1)  # prints "tracing"

At this point, the compilation cache for f contains two compiled versions of the function, one for key 0 and one for key 1. There is no need to track this cache manually. So if you call the function again with matching inputs, these cached versions of the compiled function are used, and you don't see the trace-time side effect of the print statement:

f(100, 0)  # already compiled - doesn't print anything
f(100, 1)  # also already compiled - doesn't print anything

[epignatelli]

Sorry to chip in after a while, but I am confused.

there is no way to make the hash of a tracer match its runtime value

But at staging time the real hash is not necessary because jit only has to evaluate the structure of the value (tree structure, shape, dtype, ...), or am I missing something?

Assuming that all the values have the same pytree structure, would this be a possibility (like in lax.switch)?

[jakevdp]

Staging time is the only time JAX code interacts with the Python interpreter. At runtime, all computation is happening in XLA, with no access to the python interpreter. If XLA had a dict/hashtable structure available we might be able to lower to that, but it doesn't. The only alternative would be to do something like this at runtime via a host callback, but that would be very slow.

[epignatelli]

Much clearer now, thank you @jakevdp !