Lux now allows nested AD to be done pretty seemlessly. See the announcement here. Lux also supports batched jacobians (batched_jacobian), which might be helpful depending on your use case.
1 Like
Lux now allows nested AD to be done pretty seemlessly. See the announcement here. Lux also supports batched jacobians (batched_jacobian), which might be helpful depending on your use case.