axon

axon copied to clipboard

There are cases where the input to a model is an integer type, e.g. an attention mask or token IDs. Axon currently does not respect input types, and aggressively casts to compute where necessary. For layers like embeddings, this does not make sense because they require integer inputs anyway