FilesExpand file tree

main

/

test_modeling_common_flax.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

66 lines (48 loc) · 2.69 KB

main

/

test_modeling_common_flax.py

Copy path

Top

File metadata and controls

• Code
• Blame

66 lines (48 loc) · 2.69 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

import inspect

from diffusers.utils import is_flax_available

from diffusers.utils.testing_utils import require_flax

if is_flax_available():

import jax

@require_flax

class FlaxModelTesterMixin:

def test_output(self):

init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

model = self.model_class(**init_dict)

variables = model.init(inputs_dict["prng_key"], inputs_dict["sample"])

jax.lax.stop_gradient(variables)

output = model.apply(variables, inputs_dict["sample"])

if isinstance(output, dict):

output = output.sample

self.assertIsNotNone(output)

expected_shape = inputs_dict["sample"].shape

self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")

def test_forward_with_norm_groups(self):

init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

init_dict["norm_num_groups"] = 16

init_dict["block_out_channels"] = (16, 32)

model = self.model_class(**init_dict)

variables = model.init(inputs_dict["prng_key"], inputs_dict["sample"])

jax.lax.stop_gradient(variables)

output = model.apply(variables, inputs_dict["sample"])

if isinstance(output, dict):

output = output.sample

self.assertIsNotNone(output)

expected_shape = inputs_dict["sample"].shape

self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")

def test_deprecated_kwargs(self):

has_kwarg_in_model_class = "kwargs" in inspect.signature(self.model_class.__init__).parameters

has_deprecated_kwarg = len(self.model_class._deprecated_kwargs) > 0

if has_kwarg_in_model_class and not has_deprecated_kwarg:

raise ValueError(

f"{self.model_class} has `**kwargs` in its __init__ method but has not defined any deprecated kwargs"

" under the `_deprecated_kwargs` class attribute. Make sure to either remove `**kwargs` if there are"

" no deprecated arguments or add the deprecated argument with `_deprecated_kwargs ="

" [<deprecated_argument>]`"

)

if not has_kwarg_in_model_class and has_deprecated_kwarg:

raise ValueError(

f"{self.model_class} doesn't have `**kwargs` in its __init__ method but has defined deprecated kwargs"

" under the `_deprecated_kwargs` class attribute. Make sure to either add the `**kwargs` argument to"

f" {self.model_class}.__init__ if there are deprecated arguments or remove the deprecated argument"

" from `_deprecated_kwargs = [<deprecated_argument>]`"

)