While doing some experimentation, I hit a limitation of the cache_kernel decorator that we may want to address at some point even though it is not needed right now.
Simple repro:
import numpy as np
from mujoco_warp._src.warp_util import cache_kernel
@cache_kernel
def make_kernel(n: int):
"""Returns n, simulating a kernel parameterized by n."""
return n
# Python ints: works correctly
assert make_kernel(1) == 1
assert make_kernel(6) == 6
print("Python ints: OK")
# numpy scalars: BUG — second call returns cached result from first
a = make_kernel(np.int64(1))
b = make_kernel(np.int64(6))
print(f"np.int64(1) -> {a}")
print(f"np.int64(6) -> {b}")
assert b == 6, f"Expected 6, got {b} (reused cached result from np.int64(1))"
While doing some experimentation, I hit a limitation of the
cache_kerneldecorator that we may want to address at some point even though it is not needed right now.Simple repro: