def generate()

in optimum/neuron/utils/input_generators.py [0:0]

• 50 lines of code
• 10 McCabe index (conditional complexity)

    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
        if input_name == "timestep":
            shape = [self.batch_size]
            return self.random_int_tensor(shape, max_value=999, framework=framework, dtype=int_dtype)
        elif input_name == "encoder_hidden_states":
            shape = (self.batch_size, self.sequence_length, self.text_encoder_hidden_size)
            return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
        elif input_name == "controlnet_cond":
            num_channels = getattr(
                self.normalized_config, "conditioning_channels", 3
            )  # num_channels = 3 since `do_convert_rgb=True`
            shape = (
                self.batch_size,
                num_channels,
                self.height * self.vae_scale_factor,
                self.width * self.vae_scale_factor,
            )
            return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
        elif input_name == "conditioning_scale":
            return torch.tensor([1.0])
        elif input_name == "down_block_additional_residuals":
            sample_shape = (self.batch_size, self.normalized_config.block_out_channels[0], self.height, self.width)
            sample = self.random_float_tensor(sample_shape, framework=framework, dtype=float_dtype)
            down_block_res_samples = (sample,)
            num_past_cross_attn_blocks = 0
            height = self.height
            width = self.width
            for idx, down_block_type in enumerate(self.normalized_config.down_block_types):
                res_samples = ()
                shape = (self.batch_size, self.normalized_config.block_out_channels[idx], height, width)
                for _ in range(self.normalized_config.layers_per_block):
                    res_samples += (self.random_float_tensor(shape, framework=framework, dtype=float_dtype),)
                if idx != len(self.normalized_config.down_block_types) - 1:
                    # add output of downsampler
                    num_past_cross_attn_blocks += 1
                    height = height // 2
                    width = width // 2
                    shape = (self.batch_size, self.normalized_config.block_out_channels[idx], height, width)
                    res_samples += (self.random_float_tensor(shape, framework=framework, dtype=float_dtype),)
                down_block_res_samples += res_samples
            return down_block_res_samples
        elif input_name == "mid_block_additional_residual":
            num_cross_attn_blocks = self.normalized_config.down_block_types.count("CrossAttnDownBlock2D")
            out_channels = self.normalized_config.block_out_channels[-1]
            shape = (
                self.batch_size,
                out_channels,
                self.height // 2**num_cross_attn_blocks,
                self.width // 2**num_cross_attn_blocks,
            )
            return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)