Erfan Zare Chavoshi

EasyDeL is a well-documented source library to make your training faster and more Optimized With cool Options for training and serving and supports a wide range of models that have...

Hello. I'm implementing the Mixtral models with Jax and Flax and there's a problem with the scan function at [here](https://github.com/erfanzar/EasyDeL/blob/main/lib/python/EasyDel/modules/mixtral/modelling_mixtral_flax.py#L450) and I get this error *Jax transforms and Flax models...

## Description While reviewing the codebase, I identified several areas where performance could be improved and memory usage reduced by leveraging `int1` kernels or boolean operations. These optimizations can be...

Gradient Checkpointing causes model to compute junk results (NNX)

8

### System information - **Flax, JAX, JAXlib versions**: (0.10.4) - **Python version**: (e.g., 3.10.2) ### Problem you have encountered: When using gradient checkpointing in an `flax.nnx` model (`nnx.remat`), the model...