For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

causal_conv1d_update_shape

def causal_conv1d_update_shape[dtype: DType, rank: Int](input: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec], conv_state_in: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=conv_state_in.static_spec], weight: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=weight.static_spec], bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=bias.static_spec]) -> Tuple[IndexList[rank], IndexList[rank]]

Returns the output shapes for the causal_conv1d_update op.

The update produces two tensors: the convolution output for the new token(s) and the updated convolution state.

Parameters:

dtype (DType): Element type of the input, conv state, weight, and bias tensors.
rank (Int): Tensor rank of the input and conv state, expected to be 3.

Args:

input (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec]): New input tokens with shape (batch, channels, seqlen).
conv_state_in (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=conv_state_in.static_spec]): Previous convolution state with shape (batch, channels, state_len).
weight (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=weight.static_spec]): Convolution weights with shape (channels, width).
bias (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=bias.static_spec]): Per-channel bias with shape (channels,).

Returns:

Tuple[IndexList[rank], IndexList[rank]]: A tuple (output_shape, conv_state_shape) where output_shape matches input.shape() and conv_state_shape matches conv_state_in.shape().