initial flax pndm schedular (huggingface#492)

* initial flax pndm

* fix typo

* use state

* return state

* add FlaxSchedulerOutput

* fix style

* add flax imports

* make style

* fix typos

* return created state

* make style

* add torch/flax imports

* docs

* fixed typo

* remove tensor_format

* round instead of cast

* ets is jnp array

* remove copy

Author: kashif

__init__.py

@@ -1,4 +1,5 @@
 from .utils import (
+    is_flax_available,
     is_inflect_available,
     is_onnx_available,
     is_scipy_available,
@@ -60,3 +61,8 @@
     from .pipelines import StableDiffusionOnnxPipeline
 else:
     from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
+
+if is_flax_available():
+    from .schedulers import FlaxPNDMScheduler
+else:
+    from .utils.dummy_flax_objects import *  # noqa F403

schedulers/__init__.py

@@ -12,17 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..utils import is_scipy_available
-from .scheduling_ddim import DDIMScheduler
-from .scheduling_ddpm import DDPMScheduler
-from .scheduling_karras_ve import KarrasVeScheduler
-from .scheduling_pndm import PNDMScheduler
-from .scheduling_sde_ve import ScoreSdeVeScheduler
-from .scheduling_sde_vp import ScoreSdeVpScheduler
-from .scheduling_utils import SchedulerMixin
 
+from ..utils import is_flax_available, is_scipy_available, is_torch_available
+
+
+if is_torch_available():
+    from .scheduling_ddim import DDIMScheduler
+    from .scheduling_ddpm import DDPMScheduler
+    from .scheduling_karras_ve import KarrasVeScheduler
+    from .scheduling_pndm import PNDMScheduler
+    from .scheduling_sde_ve import ScoreSdeVeScheduler
+    from .scheduling_sde_vp import ScoreSdeVpScheduler
+    from .scheduling_utils import SchedulerMixin
+else:
+    from ..utils.dummy_pt_objects import *  # noqa F403
+
+if is_flax_available():
+    from .scheduling_pndm_flax import FlaxPNDMScheduler
+else:
+    from ..utils.dummy_flax_objects import *  # noqa F403
 
 if is_scipy_available():
     from .scheduling_lms_discrete import LMSDiscreteScheduler
 else:
-    from ..utils.dummy_scipy_objects import *  # noqa F403
+    from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403