Check shape and remove deprecated APIs in scheduling_ddpm_flax.py #7703

ppham27 · 2024-04-17T19:25:06Z

model_output.shape may only have rank 1.

There are warnings related to use of random keys.

tests/schedulers/test_scheduler_flax.py: 13 warnings
  /Users/phillypham/diffusers/src/diffusers/schedulers/scheduling_ddpm_flax.py:268: FutureWarning: normal accepts a single key, but was given a key array of shape (1, 2) != (). Use jax.vmap for batching. In a future JAX version, this will be an error.
    noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype)

tests/schedulers/test_scheduler_flax.py::FlaxDDPMSchedulerTest::test_betas
  /Users/phillypham/virtualenv/diffusers/lib/python3.9/site-packages/jax/_src/random.py:731: FutureWarning: uniform accepts a single key, but was given a key array of shape (1,) != (). Use jax.vmap for batching. In a future JAX version, this will be an error.
    u = uniform(key, shape, dtype, lo, hi)  # type: ignore[arg-type]

What does this PR do?

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

HuggingFaceDocBuilderDev · 2024-04-29T22:43:44Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

yiyixuxu · 2024-04-29T22:57:35Z

src/diffusers/schedulers/scheduling_ddpm_flax.py


-        if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
+        if (
+            len(model_output.shape) > 1


why are we adding this len() > 1 check here?

Because input could have shape like x.shape == (10,). Accessing index 1 in the next line is an error in that case.

yiyixuxu · 2024-04-29T23:03:59Z

src/diffusers/schedulers/scheduling_ddpm_flax.py

@@ -264,7 +268,7 @@ def step(

        # 6. Add noise
        def random_variance():
-            split_key = jax.random.split(key, num=1)
+            split_key = jax.random.split(key, ())


why do we make this change here? num=1 would still work with jax.random.key, no?

It will be an error soon in newer versions of JAX, https://github.com/google/jax/blob/343e18fcb693c3f1b2cace56a4faea8fd3e2cadd/jax/_src/random.py#L104-L108.

In that case, I'd use:

Suggested change

split_key = jax.random.split(key, ())

split_key = jax.random.split(key, num=(1,))

That doesn't fix the batching dimension error. In any case, I have done jax.random.split(key, num=1)[0] to remove the batch dimension.

ah yes you are right!

pcuenca

Looks good! However, please note that the project dependencies are defined in

diffusers/setup.py

Line 109 in 26a7851

"jax>=0.4.1",

, and then this file is generated.

If we change this code, we need to ensure that it still works with the existing versions (can probably be confirmed with the CI tests), or upgrade the dependencies if necessary.

pcuenca · 2024-04-30T10:02:13Z

docker/diffusers-flax-cpu/Dockerfile

@@ -13,9 +13,9 @@ RUN apt update && \
                   ca-certificates \
                   libsndfile1-dev \
                   libgl1 \
-                   python3.8 \
+                   python3.10 \


The diffusers codebase is compatible with Python 3.8, and the other Dockerfiles use that version too:

diffusers/docker/diffusers-pytorch-cuda/Dockerfile

Line 16 in 26a7851

python3.8 \

.

Is it a requirement to upgrade?

No, I reverted the upgrade. But note Python 3.8 is already not supported and even Python 3.10 will lose support EOY, https://scientific-python.org/specs/spec-0000/ and https://jax.readthedocs.io/en/latest/deprecation.html.

yes, you are right, we should probably upgrade at the project level cc @yiyixuxu @sayakpaul

cc @sayakpaul here! let's upgrade?

pcuenca · 2024-04-30T10:09:19Z

src/diffusers/schedulers/scheduling_ddpm_flax.py

@@ -264,7 +268,7 @@ def step(

        # 6. Add noise
        def random_variance():
-            split_key = jax.random.split(key, num=1)
+            split_key = jax.random.split(key, ())


In that case, I'd use:

Suggested change

split_key = jax.random.split(key, ())

split_key = jax.random.split(key, num=(1,))

`model_output.shape` may only have rank 1. There are warnings related to use of random keys. ``` tests/schedulers/test_scheduler_flax.py: 13 warnings /Users/phillypham/diffusers/src/diffusers/schedulers/scheduling_ddpm_flax.py:268: FutureWarning: normal accepts a single key, but was given a key array of shape (1, 2) != (). Use jax.vmap for batching. In a future JAX version, this will be an error. noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype) tests/schedulers/test_scheduler_flax.py::FlaxDDPMSchedulerTest::test_betas /Users/phillypham/virtualenv/diffusers/lib/python3.9/site-packages/jax/_src/random.py:731: FutureWarning: uniform accepts a single key, but was given a key array of shape (1,) != (). Use jax.vmap for batching. In a future JAX version, this will be an error. u = uniform(key, shape, dtype, lo, hi) # type: ignore[arg-type] ```

pcuenca · 2024-05-08T11:57:15Z

Thanks a lot @ppham27!

…ggingface#7703) `model_output.shape` may only have rank 1. There are warnings related to use of random keys. ``` tests/schedulers/test_scheduler_flax.py: 13 warnings /Users/phillypham/diffusers/src/diffusers/schedulers/scheduling_ddpm_flax.py:268: FutureWarning: normal accepts a single key, but was given a key array of shape (1, 2) != (). Use jax.vmap for batching. In a future JAX version, this will be an error. noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype) tests/schedulers/test_scheduler_flax.py::FlaxDDPMSchedulerTest::test_betas /Users/phillypham/virtualenv/diffusers/lib/python3.9/site-packages/jax/_src/random.py:731: FutureWarning: uniform accepts a single key, but was given a key array of shape (1,) != (). Use jax.vmap for batching. In a future JAX version, this will be an error. u = uniform(key, shape, dtype, lo, hi) # type: ignore[arg-type] ```

) `model_output.shape` may only have rank 1. There are warnings related to use of random keys. ``` tests/schedulers/test_scheduler_flax.py: 13 warnings /Users/phillypham/diffusers/src/diffusers/schedulers/scheduling_ddpm_flax.py:268: FutureWarning: normal accepts a single key, but was given a key array of shape (1, 2) != (). Use jax.vmap for batching. In a future JAX version, this will be an error. noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype) tests/schedulers/test_scheduler_flax.py::FlaxDDPMSchedulerTest::test_betas /Users/phillypham/virtualenv/diffusers/lib/python3.9/site-packages/jax/_src/random.py:731: FutureWarning: uniform accepts a single key, but was given a key array of shape (1,) != (). Use jax.vmap for batching. In a future JAX version, this will be an error. u = uniform(key, shape, dtype, lo, hi) # type: ignore[arg-type] ```

ppham27 force-pushed the jax branch from 403a2c2 to 6a5597a Compare April 19, 2024 07:22

ppham27 force-pushed the jax branch from 6a5597a to d03608a Compare April 28, 2024 14:00

yiyixuxu requested a review from pcuenca April 29, 2024 22:56

yiyixuxu reviewed Apr 29, 2024

View reviewed changes

ppham27 force-pushed the jax branch 2 times, most recently from 3a655d9 to 3e6cde9 Compare April 30, 2024 01:39

pcuenca reviewed Apr 30, 2024

View reviewed changes

ppham27 force-pushed the jax branch from 3e6cde9 to 9b9e6d4 Compare April 30, 2024 13:48

pcuenca approved these changes May 8, 2024

View reviewed changes

pcuenca merged commit f29b934 into huggingface:main May 8, 2024
15 checks passed

sayakpaul mentioned this pull request May 8, 2024

upgrade to python 3.10 in the Dockerfiles #7893

Merged

	split_key = jax.random.split(key, ())
	split_key = jax.random.split(key, num=(1,))

Check shape and remove deprecated APIs in scheduling_ddpm_flax.py #7703

Check shape and remove deprecated APIs in scheduling_ddpm_flax.py #7703

Uh oh!

Conversation

ppham27 commented Apr 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Apr 29, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pcuenca left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pcuenca commented May 8, 2024

Uh oh!

Uh oh!

Uh oh!

ppham27 commented Apr 17, 2024 •

edited

Loading