DiagonalScatter

Embeds vales into the diagonal.

Abstract Signature:

DiagonalScatter(input: Tensor, src: Tensor, offset: int = 0, dim1: int = 0, dim2: int = 1)

PyTorchJAX (Core)TensorFlowFlax NNXPaxML / Praxis

PyTorch

API: torch.diagonal_scatter

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: —

Strategy: Plugin (scatter_indexer)

TensorFlow

API: tf.linalg.set_diag

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: —

Strategy: Plugin (scatter_indexer)

PaxML / Praxis

API: —

Strategy: Plugin (scatter_indexer)