Stable WarpFusion v0.33 - Reconstructed Noise Updates

Toms Zarāns

hey, i have this propblem: disabling init dummy for non-sdxl models Loading model from /content/drive/MyDrive/models/fcanimemixFcAnime_fcanimeV6.safetensors missing keys: ['betas', 'alphas_cumprod', 'alphas_cumprod_prev', 'sqrt_alphas_cumprod', 'sqrt_one_minus_alphas_cumprod', 'log_one_minus_alphas_cumprod', 'sqrt_recip_alphas_cumprod', 'sqrt_recipm1_alphas_cumprod', 'posterior_variance', 'posterior_log_variance_clipped', 'posterior_mean_coef1', 'posterior_mean_coef2', 'control_model.time_embed.0.weight', 'control_model.time_embed.0.bias', 'control_model.time_embed.2.weight', 'control_model.time_embed.2.bias', 'control_model.input_blocks.0.0.weight', 'control_model.input_blocks.0.0.bias', 'control_model.input_blocks.1.0.in_layers.0.weight', 'control_model.input_blocks.1.0.in_layers.0.bias', 'control_model.input_blocks.1.0.in_layers.2.weight', 'control_model.input_blocks.1.0.in_layers.2.bias', 'control_model.input_blocks.1.0.emb_layers.1.weight', 'control_model.input_blocks.1.0.emb_layers.1.bias', 'control_model.input_blocks.1.0.out_layers.0.weight', 'control_model.input_blocks.1.0.out_layers.0.bias', 'control_model.input_blocks.1.0.out_layers.3.weight', 'control_model.input_blocks.1.0.out_layers.3.bias', 'control_model.input_blocks.1.1.norm.weight', 'control_model.input_blocks.1.1.norm.bias', 'control_model.input_blocks.1.1.proj_in.weight', 'control_model.input_blocks.1.1.proj_in.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.norm1.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.norm1.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.norm2.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.norm2.bias', 'control_model.input_blocks.1.1.transformer_blocks.0.norm3.weight', 'control_model.input_blocks.1.1.transformer_blocks.0.norm3.bias', 'control_model.input_blocks.1.1.proj_out.weight', 'control_model.input_blocks.1.1.proj_out.bias', 'control_model.input_blocks.2.0.in_layers.0.weight', 'control_model.input_blocks.2.0.in_layers.0.bias', 'control_model.input_blocks.2.0.in_layers.2.weight', 'control_model.input_blocks.2.0.in_layers.2.bias', 'control_model.input_blocks.2.0.emb_layers.1.weight', 'control_model.input_blocks.2.0.emb_layers.1.bias', 'control_model.input_blocks.2.0.out_layers.0.weight', 'control_model.input_blocks.2.0.out_layers.0.bias', 'control_model.input_blocks.2.0.out_layers.3.weight', 'control_model.input_blocks.2.0.out_layers.3.bias', 'control_model.input_blocks.2.1.norm.weight', 'control_model.input_blocks.2.1.norm.bias', 'control_model.input_blocks.2.1.proj_in.weight', 'control_model.input_blocks.2.1.proj_in.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.norm1.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.norm1.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.norm2.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.norm2.bias', 'control_model.input_blocks.2.1.transformer_blocks.0.norm3.weight', 'control_model.input_blocks.2.1.transformer_blocks.0.norm3.bias', 'control_model.input_blocks.2.1.proj_out.weight', 'control_model.input_blocks.2.1.proj_out.bias', 'control_model.input_blocks.3.0.op.weight', 'control_model.input_blocks.3.0.op.bias', 'control_model.input_blocks.4.0.in_layers.0.weight', 'control_model.input_blocks.4.0.in_layers.0.bias', 'control_model.input_blocks.4.0.in_layers.2.weight', 'control_model.input_blocks.4.0.in_layers.2.bias', 'control_model.input_blocks.4.0.emb_layers.1.weight', 'control_model.input_blocks.4.0.emb_layers.1.bias', 'control_model.input_blocks.4.0.out_layers.0.weight', 'control_model.input_blocks.4.0.out_layers.0.bias', 'control_model.input_blocks.4.0.out_layers.3.weight', 'control_model.input_blocks.4.0.out_layers.3.bias', 'control_model.input_blocks.4.0.skip_connection.weight', 'control_model.input_blocks.4.0.skip_connection.bias', 'control_model.input_blocks.4.1.norm.weight', 'control_model.input_blocks.4.1.norm.bias', 'control_model.input_blocks.4.1.proj_in.weight', 'control_model.input_blocks.4.1.proj_in.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.norm1.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.norm1.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.norm2.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.norm2.bias', 'control_model.input_blocks.4.1.transformer_blocks.0.norm3.weight', 'control_model.input_blocks.4.1.transformer_blocks.0.norm3.bias', 'control_model.input_blocks.4.1.proj_out.weight', 'control_model.input_blocks.4.1.proj_out.bias', 'control_model.input_blocks.5.0.in_layers.0.weight', 'control_model.input_blocks.5.0.in_layers.0.bias', 'control_model.input_blocks.5.0.in_layers.2.weight', 'control_model.input_blocks.5.0.in_layers.2.bias', 'control_model.input_blocks.5.0.emb_layers.1.weight', 'control_model.input_blocks.5.0.emb_layers.1.bias', 'control_model.input_blocks.5.0.out_layers.0.weight', 'control_model.input_blocks.5.0.out_layers.0.bias', 'control_model.input_blocks.5.0.out_layers.3.weight', 'control_model.input_blocks.5.0.out_layers.3.bias', 'control_model.input_blocks.5.1.norm.weight', 'control_model.input_blocks.5.1.norm.bias', 'control_model.input_blocks.5.1.proj_in.weight', 'control_model.input_blocks.5.1.proj_in.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.norm1.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.norm1.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.norm2.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.norm2.bias', 'control_model.input_blocks.5.1.transformer_blocks.0.norm3.weight', 'control_model.input_blocks.5.1.transformer_blocks.0.norm3.bias', 'control_model.input_blocks.5.1.proj_out.weight', 'control_model.input_blocks.5.1.proj_out.bias', 'control_model.input_blocks.6.0.op.weight', 'control_model.input_blocks.6.0.op.bias', 'control_model.input_blocks.7.0.in_layers.0.weight', 'control_model.input_blocks.7.0.in_layers.0.bias', 'control_model.input_blocks.7.0.in_layers.2.weight', 'control_model.input_blocks.7.0.in_layers.2.bias', 'control_model.input_blocks.7.0.emb_layers.1.weight', 'control_model.input_blocks.7.0.emb_layers.1.bias', 'control_model.input_blocks.7.0.out_layers.0.weight', 'control_model.input_blocks.7.0.out_layers.0.bias', 'control_model.input_blocks.7.0.out_layers.3.weight', 'control_model.input_blocks.7.0.out_layers.3.bias', 'control_model.input_blocks.7.0.skip_connection.weight', 'control_model.input_blocks.7.0.skip_connection.bias', 'control_model.input_blocks.7.1.norm.weight', 'control_model.input_blocks.7.1.norm.bias', 'control_model.input_blocks.7.1.proj_in.weight', 'control_model.input_blocks.7.1.proj_in.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.norm1.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.norm1.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.norm2.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.norm2.bias', 'control_model.input_blocks.7.1.transformer_blocks.0.norm3.weight', 'control_model.input_blocks.7.1.transformer_blocks.0.norm3.bias', 'control_model.input_blocks.7.1.proj_out.weight', 'control_model.input_blocks.7.1.proj_out.bias', 'control_model.input_blocks.8.0.in_layers.0.weight', 'control_model.input_blocks.8.0.in_layers.0.bias', 'control_model.input_blocks.8.0.in_layers.2.weight', 'control_model.input_blocks.8.0.in_layers.2.bias', 'control_model.input_blocks.8.0.emb_layers.1.weight', 'control_model.input_blocks.8.0.emb_layers.1.bias', 'control_model.input_blocks.8.0.out_layers.0.weight', 'control_model.input_blocks.8.0.out_layers.0.bias', 'control_model.input_blocks.8.0.out_layers.3.weight', 'control_model.input_blocks.8.0.out_layers.3.bias', 'control_model.input_blocks.8.1.norm.weight', 'control_model.input_blocks.8.1.norm.bias', 'control_model.input_blocks.8.1.proj_in.weight', 'control_model.input_blocks.8.1.proj_in.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.norm1.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.norm1.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.norm2.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.norm2.bias', 'control_model.input_blocks.8.1.transformer_blocks.0.norm3.weight', 'control_model.input_blocks.8.1.transformer_blocks.0.norm3.bias', 'control_model.input_blocks.8.1.proj_out.weight', 'control_model.input_blocks.8.1.proj_out.bias', 'control_model.input_blocks.9.0.op.weight', 'control_model.input_blocks.9.0.op.bias', 'control_model.input_blocks.10.0.in_layers.0.weight', 'control_model.input_blocks.10.0.in_layers.0.bias', 'control_model.input_blocks.10.0.in_layers.2.weight', 'control_model.input_blocks.10.0.in_layers.2.bias', 'control_model.input_blocks.10.0.emb_layers.1.weight', 'control_model.input_blocks.10.0.emb_layers.1.bias', 'control_model.input_blocks.10.0.out_layers.0.weight', 'control_model.input_blocks.10.0.out_layers.0.bias', 'control_model.input_blocks.10.0.out_layers.3.weight', 'control_model.input_blocks.10.0.out_layers.3.bias', 'control_model.input_blocks.11.0.in_layers.0.weight', 'control_model.input_blocks.11.0.in_layers.0.bias', 'control_model.input_blocks.11.0.in_layers.2.weight', 'control_model.input_blocks.11.0.in_layers.2.bias', 'control_model.input_blocks.11.0.emb_layers.1.weight', 'control_model.input_blocks.11.0.emb_layers.1.bias', 'control_model.input_blocks.11.0.out_layers.0.weight', 'control_model.input_blocks.11.0.out_layers.0.bias', 'control_model.input_blocks.11.0.out_layers.3.weight', 'control_model.input_blocks.11.0.out_layers.3.bias', 'control_model.zero_convs.0.0.weight', 'control_model.zero_convs.0.0.bias', 'control_model.zero_convs.1.0.weight', 'control_model.zero_convs.1.0.bias', 'control_model.zero_convs.2.0.weight', 'control_model.zero_convs.2.0.bias', 'control_model.zero_convs.3.0.weight', 'control_model.zero_convs.3.0.bias', 'control_model.zero_convs.4.0.weight', 'control_model.zero_convs.4.0.bias', 'control_model.zero_convs.5.0.weight', 'control_model.zero_convs.5.0.bias', 'control_model.zero_convs.6.0.weight', 'control_model.zero_convs.6.0.bias', 'control_model.zero_convs.7.0.weight', 'control_model.zero_convs.7.0.bias', 'control_model.zero_convs.8.0.weight', 'control_model.zero_convs.8.0.bias', 'control_model.zero_convs.9.0.weight', 'control_model.zero_convs.9.0.bias', 'control_model.zero_convs.10.0.weight', 'control_model.zero_convs.10.0.bias', 'control_model.zero_convs.11.0.weight', 'control_model.zero_convs.11.0.bias', 'control_model.input_hint_block.0.weight', 'control_model.input_hint_block.0.bias', 'control_model.input_hint_block.2.weight', 'control_model.input_hint_block.2.bias', 'control_model.input_hint_block.4.weight', 'control_model.input_hint_block.4.bias', 'control_model.input_hint_block.6.weight', 'control_model.input_hint_block.6.bias', 'control_model.input_hint_block.8.weight', 'control_model.input_hint_block.8.bias', 'control_model.input_hint_block.10.weight', 'control_model.input_hint_block.10.bias', 'control_model.input_hint_block.12.weight', 'control_model.input_hint_block.12.bias', 'control_model.input_hint_block.14.weight', 'control_model.input_hint_block.14.bias', 'control_model.middle_block.0.in_layers.0.weight', 'control_model.middle_block.0.in_layers.0.bias', 'control_model.middle_block.0.in_layers.2.weight', 'control_model.middle_block.0.in_layers.2.bias', 'control_model.middle_block.0.emb_layers.1.weight', 'control_model.middle_block.0.emb_layers.1.bias', 'control_model.middle_block.0.out_layers.0.weight', 'control_model.middle_block.0.out_layers.0.bias', 'control_model.middle_block.0.out_layers.3.weight', 'control_model.middle_block.0.out_layers.3.bias', 'control_model.middle_block.1.norm.weight', 'control_model.middle_block.1.norm.bias', 'control_model.middle_block.1.proj_in.weight', 'control_model.middle_block.1.proj_in.bias', 'control_model.middle_block.1.transformer_blocks.0.attn1.to_q.weight', 'control_model.middle_block.1.transformer_blocks.0.attn1.to_k.weight', 'control_model.middle_block.1.transformer_blocks.0.attn1.to_v.weight', 'control_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.weight', 'control_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.bias', 'control_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.weight', 'control_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.bias', 'control_model.middle_block.1.transformer_blocks.0.ff.net.2.weight', 'control_model.middle_block.1.transformer_blocks.0.ff.net.2.bias', 'control_model.middle_block.1.transformer_blocks.0.attn2.to_q.weight', 'control_model.middle_block.1.transformer_blocks.0.attn2.to_k.weight', 'control_model.middle_block.1.transformer_blocks.0.attn2.to_v.weight', 'control_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.weight', 'control_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.bias', 'control_model.middle_block.1.transformer_blocks.0.norm1.weight', 'control_model.middle_block.1.transformer_blocks.0.norm1.bias', 'control_model.middle_block.1.transformer_blocks.0.norm2.weight', 'control_model.middle_block.1.transformer_blocks.0.norm2.bias', 'control_model.middle_block.1.transformer_blocks.0.norm3.weight', 'control_model.middle_block.1.transformer_blocks.0.norm3.bias', 'control_model.middle_block.1.proj_out.weight', 'control_model.middle_block.1.proj_out.bias', 'control_model.middle_block.2.in_layers.0.weight', 'control_model.middle_block.2.in_layers.0.bias', 'control_model.middle_block.2.in_layers.2.weight', 'control_model.middle_block.2.in_layers.2.bias', 'control_model.middle_block.2.emb_layers.1.weight', 'control_model.middle_block.2.emb_layers.1.bias', 'control_model.middle_block.2.out_layers.0.weight', 'control_model.middle_block.2.out_layers.0.bias', 'control_model.middle_block.2.out_layers.3.weight', 'control_model.middle_block.2.out_layers.3.bias', 'control_model.middle_block_out.0.weight', 'control_model.middle_block_out.0.bias'] 352 unexpected keys: ['model_ema.decay', 'model_ema.num_updates', 'cond_stage_model.transformer.text_model.embeddings.position_ids'] 3 --------------------------------------------------------------------------- AttributeError Traceback (most recent call last) in () 838 if 'animatediff' in model_version: 839 os.chdir(f'{root_dir}/animatediff') --> 840 from animatediff.models.motion_module import (get_motion_module, 841 TemporalTransformerBlock, TemporalTransformer3DModel) 842 os.chdir(f'{root_dir}/') 6 frames /usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py in getattr(name) 52 warnings.warn(message, DeprecationWarning, stacklevel=2) 53 return fn ---> 54 raise AttributeError(f"module {module!r} has no attribute {name!r}") 55 56 return getattr AttributeError: module 'jax.random' has no attribute 'KeyArray'