AI-Hypercomputer / AI-Hypercomputer/maxtext

`attend_dtype` not used

Open
#531 1 comment 0 reactions 1 assignee Claimed by @khatwanimohit View on GitHub
bug
Dominant language
Python
Stars
2.4k
Forks
607
Avg merge
2d 19h
Merged PRs (30d)
158

Description

[Here](https://github.com/google/maxtext/blob/5353a957594bd6cf316747cd5a327c163caca74f/MaxText/layers/embeddings.py#L102-L103) it seems that the hard-coded `bfloat16` is used instead of `attend_dtype`. Also `query` is not cast. I guess the correct behavior should be casting both `query` and `self.embedding` to `attend_dtype`?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.