Check shape and remove deprecated APIs in scheduling_ddpm_flax.py

`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]
```

Author: ppham27

docker/diffusers-flax-cpu/Dockerfile

@@ -13,9 +13,9 @@ RUN apt update && \
                    ca-certificates \
                    libsndfile1-dev \
                    libgl1 \
-                   python3.8 \
+                   python3.11 \
                    python3-pip \
-                   python3.8-venv && \
+                   python3.11-venv && \
     rm -rf /var/lib/apt/lists
 
 # make sure to use venv
@@ -27,9 +27,9 @@ ENV PATH="/opt/venv/bin:$PATH"
 RUN python3 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \
     python3 -m uv pip install --upgrade --no-cache-dir \
         clu \
-        "jax[cpu]>=0.2.16,!=0.3.2" \
-        "flax>=0.4.1" \
-        "jaxlib>=0.1.65" && \
+        "jax[cpu]>=0.4.26" \
+        "flax>=0.8.2" \
+        "jaxlib>=0.4.26" && \
     python3 -m uv pip install --no-cache-dir \
         accelerate \
         datasets \

docker/diffusers-flax-tpu/Dockerfile

@@ -13,9 +13,9 @@ RUN apt update && \
                    ca-certificates \
                    libsndfile1-dev \
                    libgl1 \
-                   python3.8 \
+                   python3.11 \
                    python3-pip \
-                   python3.8-venv && \
+                   python3.11-venv && \
     rm -rf /var/lib/apt/lists
 
 # make sure to use venv
@@ -26,12 +26,12 @@ ENV PATH="/opt/venv/bin:$PATH"
 # follow the instructions here: https://cloud.google.com/tpu/docs/run-in-container#train_a_jax_model_in_a_docker_container
 RUN python3 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \
     python3 -m pip install --no-cache-dir \
-        "jax[tpu]>=0.2.16,!=0.3.2" \
+        "jax[tpu]>=0.4.26" \
         -f https://storage.googleapis.com/jax-releases/libtpu_releases.html && \
     python3 -m uv pip install --upgrade --no-cache-dir \
         clu \
-        "flax>=0.4.1" \
-        "jaxlib>=0.1.65" && \
+        "flax>=0.8.2" \
+        "jaxlib>=0.4.26" && \
     python3 -m uv pip install --no-cache-dir \
         accelerate \
         datasets \

src/diffusers/schedulers/scheduling_ddpm_flax.py

@@ -222,9 +222,13 @@ def step(
         t = timestep
 
         if key is None:
-            key = jax.random.PRNGKey(0)
+            key = jax.random.key(0)
 
-        if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
+        if (
+            len(model_output.shape) > 1
+            and model_output.shape[1] == sample.shape[1] * 2
+            and self.config.variance_type in ["learned", "learned_range"]
+        ):
             model_output, predicted_variance = jnp.split(model_output, sample.shape[1], axis=1)
         else:
             predicted_variance = None
@@ -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, ())
             noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype)
             return (self._get_variance(state, t, predicted_variance=predicted_variance) ** 0.5) * noise