.rst .pdf

On This Page

Supported PyTorch Ops

Cerebras currently supports the following PyTorch operations.

Attention

The following list of supported PyTorch ops is very preliminary. We cannot guarantee that mixing and matching them in your models will work. Support is only provided for the way they are used in the Cerebras Reference Samples Git repository.

nn

• torch.nn.BCEWithLogitsLoss

• torch.nn.CrossEntropyLoss

• torch.nn.Dropout

• torch.nn.Embedding

• torch.nn.functional.dropout

• torch.nn.functional.gelu

• torch.nn.functional.pad

• torch.nn.functional.softmax

• torch.nn.LayerNorm

• torch.nn.Linear

• torch.nn.MSELoss

• torch.nn.NLLLoss

• torch.nn.ReLU

• torch.nn.Softmax

• torch.nn.TanH

Functional

• torch.nn.functional.log_softmax

• torch.nn.functional.relu

• torch.nn.functional.silu

Other ops

• torch.abs

• torch.all

• torch.arange

• torch.broadcast_to

• torch.cat

• torch.einsum

• torch.flatten

• torch.full

• torch.full_like

• torch.gather

• torch.log

• torch.matmul

• torch.min

• torch.ones

• torch.rsqrt

• torch.sigmoid

• torch.sum

• torch.tanh

• torch.where

• torch.zeros

• torch.zeros_like

previous

Limitations of PyTorch on Cerebras

next

Input Function Report