How to export neural network trained from JAX inside the rendering step

[ZiluLii]

This is perhaps not an issue, but a general question : )

Summary

I'm trying to use a neural network that's trained from Jax in my path integrator, is there any way to do that? Thank you!
System information:

OS: Ubuntu 20.04.6 LTS
CPU: AMD Ryzen Threadripper 1920X 12-Core Processor
GPU: NVIDIA GeForce RTX 2080 Ti
NVIDIA GeForce RTX 2080 Ti
NVIDIA GeForce RTX 2080 Ti
Python: 3.9.16 (main, Mar 8 2023, 14:00:05) [GCC 11.2.0]
NVidia driver: 530.30.02
CUDA: 11.6.55
LLVM: 15.0.7

Dr.Jit: 0.4.3
Mitsuba: 3.4.0
Is custom build? False
Compiled with: GNU 10.2.1
Variants:
scalar_rgb
scalar_spectral
cuda_ad_rgb
llvm_ad_rgb

[njroussel]

Hi @ZiluLii

Our repository has a "Discussions" tab. Feel free to use it for such questions 😄

I'm trying to use a neural network that's trained from Jax in my path integrator, is there any way to do that? Thank you!

Yes, this is possible.

You will need to disable megakernel-mode with the two following lines of code before you start rendering:

dr.set_flag(dr.JitFlag.LoopRecord, False)
dr.set_flag(dr.JitFlag.VCallRecord, False)

This existing thread elaborates a bit more on why this is necessary.