Kevin Black
Kevin Black
DenseGeneral first computes a 2-dimensional "flat shape" to initialize its kernel (see [this line](https://github.com/google/flax/blob/d9585e0a6ba6c4a4ebc93d0707add573420703df/flax/linen/linear.py#L143)) and then later reshapes the kernel to the correct shape. However, the partitioning API (e.g., `nn.with_partitioning`)...
In the README, it says to follow [this Cytron tutorial](https://www.cytron.io/tutorial/raspberry-pi-zero-usb-webcam) for installation instructions. However, the tutorial appears to be gone. Would it be possible to host it anywhere else?
I much prefer PDF.js (via [this extension](https://chromewebstore.google.com/detail/pdf-viewer/oemmndcbldboiebfnladdacbdfmadadm)) over the native Chrome PDF viewer due to the ability to go back after clicking on a reference. However, it seems that the...
### Description I'm getting a very mysterious crash on a TPU pod. It only happens during my eval step (not train step), and only happens on a multi-worker pod (not...
Hello, This is somewhat similar to #646. During training, I saved my parameters in a sharded manner (could not use aggregate because they were sharded over multiple hosts). Now I...
Hi, I'm trying to install TransformerEngine for JAX. I prefer to install cuda-toolkit via conda, and it seems like most CUDA libraries (e.g., cuDNN) and being found, but the install...
### Version nvidia-dali-cuda120==1.40.0, jax==0.4.31 ### Describe the bug. We recently discovered a problem that when we used DALI, our training curves were mysteriously worse. We were able to fix it...
### Is this a new feature, an improvement, or a change to existing functionality? Improvement ### How would you describe the priority of this feature request Should have (e.g. Adoption...
Hi @patrick-kidger, You've [discussed before](https://github.com/patrick-kidger/jaxtyping/issues/177#issuecomment-1951428395) how custom PyTrees should not call `__init__` in their unflatten function, pointing to [this section](https://jax.readthedocs.io/en/latest/pytrees.html#custom-pytrees-and-initialization) of the docs. This caused, for example, issues with `flax.struct.dataclass`...