Daniel Bershatsky

Package(s) vendoring issues

4

What is supposed way to distribute `orbax-checkpoint` и `orbax-export`? It is totally unclear for me 1. how to get notification about new releases? 2. what is 'ground truth' package(s) sources:...

Test files are not filtered properly. The issue is that `setuptools.find_package` finds packages not modules while tests are organized as a separate modules. In order to mitigate the issue, one...

Bibliography `chapter` attribute breaks bibliography entry rendering

4

### Description Field `chapter` of bibliography record breaks bibliography. It seems that kind of bibliography item (`inbook` or `article`) does not matter.  Issues stays for latest (v0.10.0) on Web...

For some reason, original CPU implementation of `maxsumexp` operation totally sucks (vectorization?). ``` -------------------------------------------------------------- Benchmark Time CPU Iterations -------------------------------------------------------------- MaxSumExpHalide/64 140812 ns 140753 ns 4985 MaxSumExpHalide/256 2290624 ns 2288591 ns...

We are interested here into T5 model only. Review low-level API and add missing ops. Review high-level (Python) API and add missing layers and ops. Add an example script for...

We are interested here into BERT and RoBERTa. - [ ] Review low-level API and add missing ops. - [ ] Review high-level (Python) API and add missing layers and...

Vendored `typeshed` is totally obsolete

1

At the moment, vendored `typeshed` package knows only Python up to 3.9 whilst the trunk has support for Python 3.12 and many third party stubs.

Operator `jax.lax.gather` consumes a lot of memory on CPU

6

### Description I'm working with large embedding layers like 3'000'000x300 or even bigger on CPU-only devices. I noticed that `jax` takes much more memory than expected. ```python import os os.environ['JAX_PLATFORM_NAME']...