You probably want something equivalent to (however you make it fast in your chosen framework): einsum('bf,fc->bfc', batched_inputs, channel_embedding) Then ca

You probably want something equivalent to (however you make it fast in your chosen framework): einsum('bf,fc->bfc', batched_inputs, channel_embedding) Then carry that info through the network and project it down at the end. It's roughly equivalent to the token embedding step in an LLM.