I don't have an environment to test this, but I believe it should be enough to get the JIT to work. This came up in https://github.com/google/jax/issues/5679 PiperOrigin-RevId: 356578451 Change-Id: I99a2aa0e87739b9edce81074fce1ca5c0bd25115
I don't have an environment to test this, but I believe it should be enough to get the JIT to work. This came up in https://github.com/google/jax/issues/5679 PiperOrigin-RevId: 356578451 Change-Id: I99a2aa0e87739b9edce81074fce1ca5c0bd25115