search

LouisDesdoigts / dLux

Differentiable optical models as parameterised neural networks in Jax using Zodiax
https://louisdesdoigts.github.io/dLux/
BSD 3-Clause "New" or "Revised" License
43 stars 6 forks source link

Pre-compile Functions with `cond`. #176

Closed Jordan-Dennis closed 1 year ago

Jordan-Dennis commented 1 year ago

Hi all, I think that we should pre-compile (using jit(inline=True)) functions containing cond. Once compiled this will make no difference to the performance (see #168), but it will improve the outside of compilation performance since cond recompiles the branches every time it is run. I believe that this is the source of the long runtimes that we have been experiencing using code. Regards Jordan.

Jordan-Dennis commented 1 year ago

A second motivation for this is that it will make the profiling more accurately reflect the performance that we are expecting.

Jordan-Dennis commented 1 year ago

So it turns out that this is not easy to do in practice so I am going to abandon this.

@benjaminpope will be familiar with the pitfalls of compiling class methods which include the requirement to mark self as static. I expected equinox to be a get out of jail free card with regards to this, but it was not. In fact because the module could be flattened it actually meant that any class with a device array leaf (all of them) could not be marked static at all :grimacing:.

benjaminpope commented 1 year ago

Yep…

——————————————— Dr Benjamin Pope (he/him) Lecturer in Astrophysics University of Queensland benjaminpope.github.io

From: Jordan Dennis @.> Sent: Wednesday, January 4, 2023 10:47:40 AM To: LouisDesdoigts/dLux @.> Cc: Benjamin Pope @.>; Mention @.> Subject: Re: [LouisDesdoigts/dLux] Pre-compile Functions with cond. (Issue #176)

So it turns out that this is not easy to do in practice so I am going to abandon this.

@benjaminpopehttps://github.com/benjaminpope will be familiar with the pitfalls of compiling class methods which include the requirement to mark self as static. I expected equinox to be a get out of jail free card with regards to this, but it was not. In fact because the module could be flattened it actually meant that any class with a device array leaf (all of them) could not be marked static at all 😬.

— Reply to this email directly, view it on GitHubhttps://github.com/LouisDesdoigts/dLux/issues/176#issuecomment-1370337026, or unsubscribehttps://github.com/notifications/unsubscribe-auth/ABN6YFOAGQBVWM55HMZGZNLWQS3BZANCNFSM6AAAAAATQEYPKQ. You are receiving this because you were mentioned.Message ID: @.***>