Memoization of function input/output

[EdwardRaff]

There appears to be some internal JAX usage of memorization in the compilation, which may be obfuscating what I'm searching for. I'm not worried about compile times in my use case, but I want to memorize the input/output pairs of some functions.

e.g., I have some function y = expensive(x) that will only ever be called with a finite number of values for x and will be called many times. I'd love a way to memoize the function expensive rather than the less ergonomic pre-calculating and looking up.

Is there any clean way to do this in Jax that still supports gradients/auto-diff?

[yashk2810]

do you want to cache the compilation or lowering (stableHLO)? We have support for both.

lowering: https://jax.readthedocs.io/en/latest/export/export.html

compilation: https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html

[EdwardRaff]

Neither I don't care about the compilation. The calculation itself is enormously expensive. e.g., I'd like to have something like:

def expensive(x):
  return x # but with a lot of work, like takes an hour to call

cache = {}
def memoized_expensive(x):
   if x in cache:
      return cache[x]
  cache[x] = expensive(x)
  return cache[x]

But in a way where I can still call jit/grad/etc on memoized_expensive

[dfm]

I don't think there's any mechanism built in to JAX that will help with this. You might be able to hack something like this manually using jax.pure_callback, although adding the global cache introduces a side effect so YMMV. Given the problem statement, it's probably better to pre-compute and use a lookup.

[EdwardRaff]

Thats unfortunate. In my case, I can pre-compute, but I could see others where you know it's finite but hard to pre-compute. Memoized function values would be a great future feature since all functions are pure! Thanks!

[jakevdp]

This may be limited by hardware capabilities: for example, I'm not sure there's any good mechanism for memoized computation on GPU without using a host-side cache, in which case a pure_callback solution would be about the the best you could do.