Francesco Zanetta

I am using TF2.7 and TFP0.15. If I pass the `save_traces=False` argument to `model.save()` it works. ```python import numpy as np import tensorflow as tf import tensorflow_probability as tfp from...

> This is very unsafe, see for example: https://simonbyrne.github.io/notes/fastmath/ Thanks a lot for the link @dcherian . I wish there were more blog posts like this. So with the fastmath...

> That's a cool example, thanks @frazane . > > ...but unfortunately it looks `fastmath` this does: > > > Allow optimizations for floating-point arithmetic that assume that arguments and...

@dcherian @max-sixty I've looked into the available fastmath flags (https://llvm.org/docs/LangRef.html#fast-math-flags) and it looks like the "no NaNs" and "no Infs" flags can be turned off. ```python fastmath_flags = { "nnan":...

> Could we confirm that nans are handled correctly? I guess a more careful inspection of the generated asm might reveal whether nan checks are actually done. But the fact...

> > I guess a more careful inspection of the generated asm might reveal whether nan checks are actually done. But the fact that the `np.allclose` evaluates to True is...

So, appartently it's much easier than I imagined. Although it's not well documented, `guvectorize` does take `fastmath` as a keyword argument. So it boils down to: ```python if os.getenv("NUMBAGG_FASTMATH", 'False').lower()...

> Seems like this is intentional behavior as there's a [line of code](https://github.com/google/flax/blob/main/flax/experimental/nnx/nnx/module.py#L231) that catches whether the input is a `jax.Array` or not. Any thoughts @cgarciae? It's intentional that `jax.Array`...

It would be amazing to have a built-in IDE debugger (for example in VSCode) be able to debug snakemake python files. I used snakemake for a project, and it's a...

Hi Cristian, love the work you're doing with this library 🔥 Quick question related to this issue: is currently possible to access `Variable` metadata from within a `nnx.Module` (e.g. inside...