partial(fn, **kwargs) that defines hash and eq - to help with jit caches. This is particularly common in evaluators when one has many evaluator instances that run on difference slices of data. Example: ``` f1 = _CacheablePartial(fn, a=1) jax.jit(f1)(...) jax.jit(_CacheablePart
| 80 | |
| 81 | @dataclasses.dataclass(frozen=True, eq=True) |
| 82 | class _CacheablePartial: |
| 83 | """partial(fn, **kwargs) that defines hash and eq - to help with jit caches. |
| 84 | |
| 85 | This is particularly common in evaluators when one has many evaluator |
| 86 | instances that run on difference slices of data. |
| 87 | |
| 88 | Example: |
| 89 | |
| 90 | ``` |
| 91 | f1 = _CacheablePartial(fn, a=1) |
| 92 | jax.jit(f1)(...) |
| 93 | jax.jit(_CacheablePartial(fn, a=1))(...) # fn won't be retraced. |
| 94 | del f1 |
| 95 | jax.jit(_CacheablePartial(fn, a=1))(...) # fn will be retraced. |
| 96 | ``` |
| 97 | """ |
| 98 | fn: Callable[..., Any] |
| 99 | kwargs: flax.core.FrozenDict |
| 100 | |
| 101 | def __call__(self, *args, **kwargs): |
| 102 | return functools.partial(self.fn, **self.kwargs)(*args, **kwargs) |