Version 0.6.5

[IvyZX]

What's Changed

• Internal change by @copybara-service in Internal change #2826
• Cache .mypy_cache on CI by @cgarciae in Cache .mypy_cache on CI #2819
• improve nn.map_variables docstrings by @cgarciae in improve nn.map_variables docstrings #2818
• Render non-array variables in tabulate by @cgarciae in Render non-array variables in tabulate #2769
• Improve nn.map_variables example by @cgarciae in Improve nn.map_variables example #2831
• Introduce relaxed naming policy behind a feature flag. by @levskaya in Introduce relaxed naming policy behind a feature flag. #2833
• Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code. by @copybara-service in Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code. #2840
• import initializers module directly by @chiamp in import initializers module directly #2830
• Re-box axis metadata in DenseGeneral's init functions by @IvyZX in Re-box axis metadata in DenseGeneral's init functions #2843
• Update jax dependency by @jheek in Update jax dependency #2827
• replaced zeros and ones initializers with zeros_init and ones_init builder functions by @chiamp in replaced zeros and ones initializers with zeros_init and ones_init builder functions #2815
• Fixing the punctuation in the README by @arcAman07 in Fixing the punctuation in the README #2811
• Implement lazy init by @jheek in Implement lazy init #2816
• improve nn.checkpoint docs by @cgarciae in improve nn.checkpoint docs #2837
• Flax 0.6.5 point release by @levskaya in Flax 0.6.5 point release #2856

Full Changelog: v0.6.4...v0.6.5

---

This discussion was created from the release Version 0.6.5.

[hkjeon13]

This version has the problem that the FlaxModel from the huggingface can't use the GPUs on the Colab Notebook. I tried many ways based on the error message, but any solution didn't work except installing the previous version <=0.6.4.

Here is my Colab Notebook script with GPU (just for the explanation).

#Flax-trainer is my PyPI module for the training, but it has no dependencies except jax, jaxlib, transformers, datasets, evaluate, and flax.
!pip install git+https://github.com/hkjeon13/flax-trainer.git

from flax_trainer import FlaxTrainerForSequenceClassification
from transformers import FlaxAutoModelForSequenceClassification, TrainingArguments, AutoTokenizer
from datasets import load_dataset


tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
dataset = load_dataset("glue", "cola", split="train")
print(dataset)
def example_function(examples):
    tokenized_inputs = tokenizer(
        examples["sentence"], 
        padding='max_length', 
        max_length=256, 
        truncation=True
    )
    if "label" in examples:
        tokenized_inputs["labels"] = examples["label"]
    return tokenized_inputs

dataset = dataset.map(example_function, batched=True, remove_columns=dataset.column_names)

model = FlaxAutoModelForSequenceClassification.from_pretrained("bert-base-cased", from_pt=True)
args = TrainingArguments("runs")
trainer = FlaxTrainerForSequenceClassification(
    model=model,
    args = args,
    train_dataset=dataset
)

trainer.train()

I tested on the GPU env, but flax (==0.6.5) didn't recognize the GPU resource.
I also do the updates on jax, jaxlib for the flax(==0.6.5), but it occurs dependencies problem.

[levskaya]

Hi - could you give any more information on what the error messages were?

There's very little in flax itself that would affect GPU compatibility.

I'm guessing there might be an issue with recent jax / jaxlib that we require with new minimum version requirements?

We can try to track down what's happening, but more info on the errors would be helpful.

[cgarciae]

I managed to get transformers working on colab (cpu) by updating jax before installing flax + transformers:

! pip install -U jax jaxlib
! pip install flax transformers[flax]

[hkjeon13]

Hi, Thanks for your concern. It is the best way for the error sharing my Colab Notebook.
Here is the notebook link on github gist: https://gist.github.com/hkjeon13/20bbf42637a75a1d29f5ea9e5276b2a3
In this notebook, you could find the error "WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)" even on the GPU environment.

[andsteing]

Can you try updating the PIP packages like this?

jax[cuda11_cudnn82]>=0.4.2
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

(from https://github.com/google-research/vision_transformer/blob/main/vit_jax/requirements.txt )

[hkjeon13]

Can you try updating the PIP packages like this?

jax[cuda11_cudnn82]>=0.4.2
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

(from https://github.com/google-research/vision_transformer/blob/main/vit_jax/requirements.txt )

I tried this and got a successful training result. Thank you for the suggestion:)

[wookayin]

Why was 0.6.5 version yanked? I cannot find a relevant thread anywhere.

https://pypi.org/project/flax/0.6.5/