Creating a custom random module that abstracts Jax's verbose process can be a useful way to simplify random number generation and make your code more readable.

Creating a custom dropout layer that abstracts the complicated dropout implementation in Flax/Jax can make it easier to use dropout in various Jax setups.

With the latest jax=='0.4.25' and jaxlib=='0.4.25' I get:

import nanodl
# AttributeError: module 'jax.random' has no attribute 'KeyArray'

Hi @HMUNACHI ,

Your dataloader code is heavily referred to the jax-dataloader repository, which is totally okay. But could you kindly cite the repository for reference?

Thank you!

Norms across most Attention based implementations have epsilon value set to dropout probability

Implementation of IJEPA

ReadTheDocs is failing.

I'm a Cyber Security researcher and developer of PackjGuard [1] to address open-source software supply chain attacks.

Issue

During my research, I detected a deleted package in this repository.

Details

Specifically, the package nanodl mentioned in file README at line 43 does not exist on the public PyPI registry. A bad actor can hijack this package to propagate malicious code.

Impact

Not only your apps/services using https://github.com/HMUNACHI/nanodl repo code are vulnerable to this attack, but the users of your open-source Github repo could also fall victim.

You could read more about such attacks here: https://medium.com/@alex.birsan/dependency-confusion-4a5d60fec610

Remediation

Please highlight this in file README and register a placeholder package for nanodl on public PyPI soon to remediate.

To automatically fix such issues in future, please install PackjGuard Github app [1].

Thanks!

1. PackjGuard is a Github app that monitors your repos 24x7, detects vulnerable/malicious/risky open-source dependencies, and creates pull requests for auto remediation: https://github.com/marketplace/packjguard

Hey, great job with nanodl!

I was just looking through the code and noticed that when in Lambda's Trainer the gradients are not being averaged across devices here:

Not sure if this is happening elsewhere but usually to keep the weights in sync you apply a jax.lax.pmean over the gradients before passing them to apply_gradients, e.g.

grads = jax.lax.pmean(grads, axis_name='devices')