text: 
  
  TITLE
  [Flax] from_pretrained does not consider the passed dtype
  
  COMMENTS
  6
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  ## Environment info
  
  When loading a flax model with `from_pretrained` the type argument is not used. The weights are initialized with the dtype of saved weights.
  
  So if you do 
  ```python
  model = FlaxGPT2ForCausalLM.from_pretrained("gpt2", dtype=jnp.dtype("bfloat16"))
  
  # check the dtype of one of the params 
  model.params["transformer"]["wpe"]["embedding"].dtype
  => dtype("float32")
  ```
  
  We should probably cast the weights to `self.dtype`.
  
  As a workaround for `bf16`,  one could manually cast the weighs with
  ```
  def to_bf16(t):
      return jax.tree_map(lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t)
  
  model.params = to_bf16(model.params)
  ```
  
  cc @patrickvonplaten 