ProjectLastDim

Linear projection on the last dim of the input Tensor using Einsum.

Abstract Signature:

ProjectLastDim(inputs: Tensor, weight: Tensor)

PyTorchJAX (Core)PaxML / Praxis

PyTorch

API: torch.nn.functional.linear

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: jax.numpy.matmul

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: praxis.layers.project_last_dim

Strategy: Direct Mapping

Official Docs ↗